Update app.py

app.py

@@ -1,128 +1,65 @@
-import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
 
-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",
-    live=True,
-    title="Airavata LLMs Chatbot",
-    description="Ask me anything, and I'll generate a response!",
-    theme="light",
-)
-
-iface.launch()
-
-
-
-
-
-
-
-
-
-
-
-
-# import torch
-# from transformers import AutoTokenizer, AutoModelForCausalLM
-# import gradio as gr
-
-# model_name = "ai4bharat/Airavata"
-# tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
-# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
-
-# SYSTEM_PROMPT = """<s>[INST] <<SYS>>
-# नमस्कार! आप अब कृषि विशेषज्ञता बॉट के साथ इंटरैक्ट कर रहे हैं—एक उन्नत AI जो कृषि क्षेत्र में विशेषज्ञता प्रदान करने के लिए डिज़ाइन किया गया है।
-
-# कृपया ध्यान दें कि यह बॉट केवल हिंदी में जवाब देगा। इसकी क्षमताएँ शामिल हैं:
-
-# 1. आधुनिक फसल प्रबंधन तकनीकों में गहरा ज्ञान।
-# 2. कृषि में कीट और रोग नियंत्रण के लिए प्रभावी रणनीतियाँ।
-# 3. मृदा स्वास्थ्य का सुधारने और पुनर्निर्माण के लिए विशेषज्ञता।
-# 4. सतत और प्रेसिजन खेती के अभ्यासों का ज्ञान।
-# 5. सिंचाई और जल प्रबंधन के लिए सर्वोत्तम अभ्यासों के लिए सुझाव।
-# 6. रणनीतिक फसल चक्रण और इंटरक्रॉपिंग विधियों पर मार्गदर्शन।
-# 7. नवीनतम कृषि प्रौद्योगिकियों और नवाचारों की जानकारी।
-# 8. विशेष फसलों, जलवायु, और क्षेत्रों के लिए विशेषज्ञ सलाह।
-
-# कृपया पेशेवर रूप से बराबरी बनाए रखें और सुनिश्चित करें कि आपके जवाब सही और मूल्यवान हैं। उपयोगकर्ताओं से आगे की स्पष्टीकरण के लिए पूछने के लिए प्रोत्साहित करें।
-
-# आपका प्रमुख लक्ष्य है यह है कि आप कृषि क्षेत्र में उपयुक्त ज्ञान प्रदान करें। आपके ज्ञान का धन्यवाद।
-# <</SYS>>
-# """
-
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# def create_prompt_with_chat_format(messages, bos="<s>", eos="</s>", add_bos=True, system_prompt="System: "):
-#     formatted_text = ""
-#     for message in messages:
-#         if message["role"] == "system":
-#             formatted_text += system_prompt + message["content"] + "\n"
-#         elif message["role"] == "user":
-#             if isinstance(message["content"], list):
-#                 formatted_text += "\n" + "\n".join(message["content"]) + "\n"
-#             else:
-#                 formatted_text += "\n" + message["content"] + "\n"
-#         elif message["role"] == "assistant":
-#             if isinstance(message["content"], list):
-#                 formatted_text += "\n" + "\n".join(message["content"]).strip() + eos + "\n"
-#             else:
-#                 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
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-# def inference(input_prompts, model, tokenizer, system_prompt="System: "):
-#     output_texts = []
-#     model = model.to(device)  # Move the model to the same device as the input data
-#     for input_prompt in input_prompts:
-#         formatted_query = create_prompt_with_chat_format([{"role": "user", "content": input_prompt}], add_bos=False, system_prompt=system_prompt)
-#         encodings = tokenizer(formatted_query, padding=True, return_tensors="pt")
-#         encodings = {key: value.to(device) for key, value in encodings.items()}  # Move input data to the same device as the model
+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
 
-#         with torch.no_grad():
-#             outputs = model.generate(encodings["input_ids"], do_sample=False, max_length=250)
 
-#         output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-#         output_texts.append(output_text[len(input_prompt):])
-#     return output_texts
+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)
 
-# examples = [
-#     ["मुझे अपने करियर के बारे में सुझाव दो", "मैं कैसे अध्ययन कर सकता हूँ?"],
-#     ["कृपया मुझे एक कहानी सुनाएं", "ताजमहल के बारे में कुछ बताएं"],
-#     ["मेरा नाम क्या है?", "आपका पसंदीदा फिल्म कौन सी है?"],
-# ]
+    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 get_llama_response(message: str, history: list, system_prompt=SYSTEM_PROMPT) -> str:
-#     formatted_history = [{"role": "user", "content": hist} for hist in history]
-#     formatted_message = {"role": "user", "content": message}
+model_name = "ai4bharat/Airavata"
 
-#     formatted_query = create_prompt_with_chat_format(formatted_history + [formatted_message], add_bos=False, system_prompt=system_prompt)
-#     response = inference([formatted_query], model, tokenizer)
+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("Chatbot:", response[0].strip())
-#     return response[0].strip()
+examples= [
+    "मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं।",
+    "मैं अपने समय प्रबंधन कौशल को कैसे सुधार सकता हूँ? मुझे पांच बिंदु बताएं और उनका वर्णन करें।",
+]
+# outputs = inference(input_prompts, model, tokenizer)
+# print(outputs)
 
+gr.ChatInterface(fn=inference,
+                examples = examples,
+                title = "CAMAI ChatBot").launch()
 
-# gr.ChatInterface(fn=get_llama_response).launch()