Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 tokenizer = AutoTokenizer.from_pretrained("ai4bharat/Airavata")
 model = AutoModelForCausalLM.from_pretrained("ai4bharat/Airavata")
 
-def generate_response(prompt):
+def chat_interface(prompt):
     # Tokenize input prompt and generate response
     inputs = tokenizer(prompt, return_tensors="pt", max_length=256, truncation=True)
     outputs = model.generate(**inputs)
@@ -13,23 +13,22 @@ def generate_response(prompt):
     
     return response
 
-# Define Gradio interface
-iface = gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(),
-    outputs=gr.Textbox(),
-    live=True,
-    title="CAMAI",
-    description="Enter a prompt to generate text.",
+# Define Gradio Chat Interface
+iface = gr.ChatInterface(
+    chat_model=chat_interface,
+    title="GPT-2 Chat Interface",
+    inputs=["text"],
+    outputs=["text"],
+    examples = [
+    ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं।"],
+    ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं और उनका वर्णन करें।"],
+    ],
 )
 
-# Launch Gradio interface
+# Launch Gradio Chat Interface
 iface.launch()
 
 
-
-
-
 # import torch
 # from transformers import AutoTokenizer, AutoModelForCausalLM
 # import gradio as gr
@@ -78,24 +77,24 @@ iface.launch()
 #     return output_texts
 
 
-model_name = "ai4bharat/Airavata"
+# model_name = "ai4bharat/Airavata"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
-tokenizer.pad_token = tokenizer.eos_token
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
-print(f"Loading model: {model_name}")
+# tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
+# tokenizer.pad_token = tokenizer.eos_token
+# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
+# print(f"Loading model: {model_name}")
 
-examples = [
-    ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं।"],
-    ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं और उनका वर्णन करें।"],
-]
+# examples = [
+#     ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं।"],
+#     ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं और उनका वर्णन करें।"],
+# ]
 
-def chat_interface(input_prompts):
-    outputs = inference(input_prompts, model, tokenizer)
-    return outputs
+# def chat_interface(input_prompts):
+#     outputs = inference(input_prompts, model, tokenizer)
+#     return outputs
 
-gr.Interface(fn=chat_interface,
-             inputs="text",
-             outputs="text",
-             examples=examples,
-             title="CAMAI ChatBot").launch()
+# gr.Interface(fn=chat_interface,
+#              inputs="text",
+#              outputs="text",
+#              examples=examples,
+#              title="CAMAI ChatBot").launch()