Update app.py

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Initialize the model
+# Initialize the model and tokenizer
 model_name = "sarvamai/sarvam-2b-v0.5"
-pipe = pipeline("text-generation", model=model_name, device=-1)  # Set device to -1 for CPU
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
 # Supported languages
 LANGUAGES = ["English", "Bengali", "Gujarati", "Hindi", "Kannada", "Malayalam", "Marathi", "Oriya", "Punjabi", "Tamil", "Telugu"]
@@ -15,8 +17,20 @@ def chatbot(message, history, language):
         prompt += f"Human: {human}\nAI: {ai}\n"
     prompt += f"Human: {message}\nAI:"
 
-    # Generate response
-    response = pipe(prompt, max_new_tokens=100, temperature=0.7, repetition_penalty=1.1)[0]['generated_text']
+    # Tokenize and generate
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=100,
+            temperature=0.7,
+            repetition_penalty=1.1,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     # Extract only the AI's response
     ai_response = response.split("AI:")[-1].strip()
@@ -32,10 +46,20 @@ iface = gr.ChatInterface(
     title="Multilingual Indian Chatbot",
     description="Chat in multiple Indian languages using the sarvam-2b model.",
     examples=[
-        ["नमस्ते, आप कैसे हैं?", "Hindi"]
+        ["Hello, how are you?", "English"],
+        ["नमस्ते, आप कैसे हैं?", "Hindi"],
+        ["வணக்கம், எப்படி இருக்கிறீர்கள்?", "Tamil"],
+        ["ନମସ୍କାର, ଆପଣ କେମିତି ଅଛନ୍ତି?", "Oriya"],
+        ["નમસ્તે, તમે કેમ છો?", "Gujarati"],
+        ["নমস্কার, আপনি কেমন আছেন?", "Bengali"],
+        ["ನಮಸ್ಕಾರ, ನೀವು ಹೇಗಿದ್ದೀರಿ?", "Kannada"],
+        ["നമസ്കാരം, സുഖമാണോ?", "Malayalam"],
+        ["नमस्कार, तुम्ही कसे आहात?", "Marathi"],
+        ["ਸਤ ਸ੍ਰੀ ਅਕਾਲ, ਤੁਸੀਂ ਕਿਵੇਂ ਹੋ?", "Punjabi"],
+        ["నమస్కారం, మీరు ఎలా ఉన్నారు?", "Telugu"]
     ],
     theme="soft"
 )
 
 # Launch the interface
-iface.launch()
+iface.launch()