updated with interface

app.py

@@ -1,100 +1,65 @@
-import gradio as gr
+
+
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+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)
+
+
+def create_prompt_with_chat_format(messages, bos="<s>", eos="</s>", add_bos=True):
+    formatted_text = ""
+    for message in messages:
+        if message["role"] == "system":
+            formatted_text += "<|system|>\n" + message["content"] + "\n"
+        elif message["role"] == "user":
+            formatted_text += "<|user|>\n" + message["content"] + "\n"
+        elif message["role"] == "assistant":
+            formatted_text += "<|assistant|>\n" + message["content"].strip() + eos + "\n"
+        else:
+            raise ValueError(
+                "Tulu chat template only supports 'system', 'user' and 'assistant' roles. Invalid role: {}.".format(
+                    message["role"]
+                )
+            )
+    formatted_text += "<|assistant|>\n"
+    formatted_text = bos + formatted_text if add_bos else formatted_text
+    return formatted_text
+
+
+def inference(input_prompts, model, tokenizer):
+    input_prompts = [
+        create_prompt_with_chat_format([{"role": "user", "content": input_prompt}], add_bos=False)
+        for input_prompt in input_prompts
+    ]
+
+    encodings = tokenizer(input_prompts, padding=True, return_tensors="pt")
+    encodings = encodings.to(device)
+
+    with torch.inference_mode():  # Add missing import statement for torch.inference_mode()
+        outputs = model.generate(encodings.input_ids, do_sample=False, max_new_tokens=250)
+
+    output_texts = tokenizer.batch_decode(outputs.detach(), skip_special_tokens=True)
+
+    input_prompts = [
+        tokenizer.decode(tokenizer.encode(input_prompt), skip_special_tokens=True) for input_prompt in input_prompts
+    ]
+    output_texts = [output_text[len(input_prompt) :] for input_prompt, output_text in zip(input_prompts, output_texts)]
+    return output_texts
+
+
+def chat_interface(input_prompts):
+    outputs = inference(input_prompts, model, tokenizer)
+    return outputs
+
+
+inputs = gr.inputs.Textbox(lines=2, label="User Input")
+outputs = gr.outputs.Textbox(label="Assistant Response")
 
-# Load model and tokenizer directly
-tokenizer = AutoTokenizer.from_pretrained("ai4bharat/Airavata")
-model = AutoModelForCausalLM.from_pretrained("ai4bharat/Airavata")
-
-def chat_interface(user_input, assistant_input):
-    # Concatenate the user and assistant inputs to simulate a chat conversation
-    chat_history = f"{assistant_input} User: {user_input}"
-
-    # Tokenize the chat history and generate the response
-    inputs = tokenizer(chat_history, return_tensors="pt", max_length=256, truncation=True)
-    outputs = model.generate(**inputs)
-    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    
-    return response, chat_history
-
-# Define Gradio Chat Interface
-iface = gr.ChatInterface(
-    chat_model=chat_interface,
-    title="GPT-2 Chat Interface",
-    inputs=["text", "text"],
-    outputs=["text", "text"],
-)
-
-# Launch Gradio Chat Interface
-iface.launch()
-
-
-
-# import torch
-# from transformers import AutoTokenizer, AutoModelForCausalLM
-# import gradio as gr
-
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-# def create_prompt_with_chat_format(messages, bos="<s>", eos="</s>", add_bos=True):
-#     formatted_text = ""
-#     for message in messages:
-#         if message["role"] == "system":
-#             formatted_text += "<|system|>\n" + message["content"] + "\n"
-#         elif message["role"] == "user":
-#             formatted_text += "<|user|>\n" + message["content"] + "\n"
-#         elif message["role"] == "assistant":
-#             formatted_text += "<|assistant|>\n" + message["content"].strip() + eos + "\n"
-#         else:
-#             raise ValueError(
-#                 "Tulu chat template only supports 'system', 'user' and 'assistant' roles. Invalid role: {}.".format(
-#                     message["role"]
-#                 )
-#             )
-#     formatted_text += "<|assistant|>\n"
-#     formatted_text = bos + formatted_text if add_bos else formatted_text
-#     return formatted_text
-
-
-# def inference(input_prompts, model, tokenizer):
-#     input_prompts = [
-#         create_prompt_with_chat_format([{"role": "user", "content": input_prompt}], add_bos=False)
-#         for input_prompt in input_prompts
-#     ]
-
-#     encodings = tokenizer(input_prompts, padding=True, return_tensors="pt")
-#     encodings = encodings.to(device)
-
-#     with torch.inference_mode():
-#         outputs = model.generate(encodings.input_ids, do_sample=False, max_new_tokens=250)
-
-#     output_texts = tokenizer.batch_decode(outputs.detach(), skip_special_tokens=True)
-
-#     input_prompts = [
-#         tokenizer.decode(tokenizer.encode(input_prompt), skip_special_tokens=True) for input_prompt in input_prompts
-#     ]
-#     output_texts = [output_text[len(input_prompt) :] for input_prompt, output_text in zip(input_prompts, output_texts)]
-#     return output_texts
-
-
-# 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}")
-
-# examples = [
-#     ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं।"],
-#     ["मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं और उनका वर्णन करें।"],
-# ]
-
-# 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=inputs, outputs=outputs, title="Chat Interface").launch()