Update app.py

app.py

@@ -1,40 +1,34 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM,pipeline
+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)
 
-model = "ai4bharat/Airavata"
+SYSTEM_PROMPT = """<s>[INST] <<SYS>>
+नमस्कार! आप अब कृषि विशेषज्ञता बॉट के साथ इंटरैक्ट कर रहे हैं—एक उन्नत AI जो कृषि क्षेत्र में विशेषज्ञता प्रदान करने के लिए डिज़ाइन किया गया है।
 
-tokenizer = AutoTokenizer.from_pretrained(model, padding_side="left")
-# tokenizer.pad_token = tokenizer.eos_token
-# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
+कृपया ध्यान दें कि यह बॉट केवल हिंदी में जवाब देगा। इसकी क्षमताएँ शामिल हैं:
 
-llama_pipeline = pipeline(
-    "text-generation",
-    model=model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-SYSTEM_PROMPT = """<s>[INST] <<SYS>>
-You are a helpful bot. Your answers are clear and concise.
+1. आधुनिक फसल प्रबंधन तकनीकों में गहरा ज्ञान।
+2. कृषि में कीट और रोग नियंत्रण के लिए प्रभावी रणनीतियाँ।
+3. मृदा स्वास्थ्य का सुधारने और पुनर्निर्माण के लिए विशेषज्ञता।
+4. सतत और प्रेसिजन खेती के अभ्यासों का ज्ञान।
+5. सिंचाई और जल प्रबंधन के लिए सर्वोत्तम अभ्यासों के लिए सुझाव।
+6. रणनीतिक फसल चक्रण और इंटरक्रॉपिंग विधियों पर मार्गदर्शन।
+7. नवीनतम कृषि प्रौद्योगिकियों और नवाचारों की जानकारी।
+8. विशेष फसलों, जलवायु, और क्षेत्रों के लिए विशेषज्ञ सलाह।
+
+कृपया पेशेवर रूप से बराबरी बनाए रखें और सुनिश्चित करें कि आपके जवाब सही और मूल्यवान हैं। उपयोगकर्ताओं से आगे की स्पष्टीकरण के लिए पूछने के लिए प्रोत्साहित करें।
+
+आपका प्रमुख लक्ष्य है यह है कि आप कृषि क्षेत्र में उपयुक्त ज्ञान प्रदान करें। आपके ज्ञान का धन्यवाद।
 <</SYS>>
 
 """
 
 # Formatting function for message and history
 def format_message(message: str, history: list, memory_limit: int = 3) -> str:
-    """
-    Formats the message and history for the Llama model.
-
-    Parameters:
-        message (str): Current message to send.
-        history (list): Past conversation history.
-        memory_limit (int): Limit on how many past interactions to consider.
-
-    Returns:
-        str: Formatted message string
-    """
-    # always keep len(history) <= memory_limit
     if len(history) > memory_limit:
         history = history[-memory_limit:]
 
@@ -43,45 +37,30 @@ def format_message(message: str, history: list, memory_limit: int = 3) -> str:
 
     formatted_message = SYSTEM_PROMPT + f"{history[0][0]} [/INST] {history[0][1]} </s>"
 
-    # Handle conversation history
     for user_msg, model_answer in history[1:]:
         formatted_message += f"<s>[INST] {user_msg} [/INST] {model_answer} </s>"
 
-    # Handle the current message
     formatted_message += f"<s>[INST] {message} [/INST]"
 
     return formatted_message
 
-# Generate a response from the Llama model
-def get_llama_response(message: str, history: list) -> str:
-    """
-    Generates a conversational response from the Llama model.
-
-    Parameters:
-        message (str): User's input message.
-        history (list): Past conversation history.
+def inference(input_prompts, model, tokenizer):
+    input_prompts = [
+        tokenizer.encode(input_prompt, return_tensors="pt", max_length=1024, truncation=True)
+        for input_prompt in input_prompts
+    ]
 
-    Returns:
-        str: Generated response from the Llama model.
-    """
-    query = format_message(message, history)
-    response = ""
-
-    sequences = llama_pipeline(
-        query,
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=tokenizer.eos_token_id,
-        max_length=1024,
-    )
+    with torch.inference_mode():
+        outputs = model.generate(input_prompts[0], do_sample=True, top_k=10, max_length=1024)
 
-    generated_text = sequences[0]['generated_text']
-    response = generated_text[len(query):]  # Remove the prompt from the output
+    output_texts = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return output_texts
 
+def get_llama_response(message: str, history: list) -> str:
+    query = format_message(message, history)
+    response = inference([query], model, tokenizer)
+    
     print("Chatbot:", response.strip())
     return response.strip()
 
-
-
 gr.ChatInterface(get_llama_response).launch()