Update app.py

app.py

@@ -1,65 +1,98 @@
 import gradio as gr
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-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 += "\n" + message["content"] + "\n"
-        elif message["role"] == "user":
-            formatted_text += "\n" + message["content"] + "\n"
-        elif message["role"] == "assistant":
-            formatted_text += "\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 += "\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.no_grad():
-        outputs = model.generate(encodings.input_ids, do_sample=False, max_length=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)
-examples = [
-    ["मुझे अपने करियर के बारे में सुझाव दो", "मैं कैसे अध्ययन कर सकता हूँ?"],
-    ["कृपया मुझे एक कहानी सुनाएं", "ताजमहल के बारे में कुछ बताएं"],
-    ["मेरा नाम क्या है?", "आपका पसंदीदा फिल्म कौन सी है?"],
-]
-
-iface = gr.Chat(
-    model_fn=lambda input_prompts: inference(input_prompts, model, tokenizer),
-    inputs=["text"],
+tokenizer = AutoTokenizer.from_pretrained("ai4bharat/Airavata")
+model = AutoModelForCausalLM.from_pretrained("ai4bharat/Airavata")
+
+def generate_response(prompt):
+    input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=50)
+    output_ids = model.generate(input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2)
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return response
+
+iface = gr.Interface(
+    fn=generate_response,
+    inputs="text",
     outputs="text",
-    examples=examples,
-    title="Airavata Chatbot",
-    theme="light",  # Optional: Set a light theme
+    live=True,
+    title="Airavata LLMs Chatbot",
+    description="Ask me anything, and I'll generate a response!",
+    theme="light",
 )
 
 iface.launch()
+
+
+
+
+
+
+
+
+
+
+# import gradio as gr
+# import torch
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# 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 += "\n" + message["content"] + "\n"
+#         elif message["role"] == "user":
+#             formatted_text += "\n" + message["content"] + "\n"
+#         elif message["role"] == "assistant":
+#             formatted_text += "\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 += "\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.no_grad():
+#         outputs = model.generate(encodings.input_ids, do_sample=False, max_length=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)
+# examples = [
+#     ["मुझे अपने करियर के बारे में सुझाव दो", "मैं कैसे अध्ययन कर सकता हूँ?"],
+#     ["कृपया मुझे एक कहानी सुनाएं", "ताजमहल के बारे में कुछ बताएं"],
+#     ["मेरा नाम क्या है?", "आपका पसंदीदा फिल्म कौन सी है?"],
+# ]
+
+# iface = gr.Chat(
+#     model_fn=lambda input_prompts: inference(input_prompts, model, tokenizer),
+#     inputs=["text"],
+#     outputs="text",
+#     examples=examples,
+#     title="Airavata Chatbot",
+#     theme="light",  # Optional: Set a light theme
+# )
+
+# iface.launch()