Update app.py

app.py

@@ -1,25 +1,40 @@
 import spaces
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Initialize the pipeline
-pipe = pipeline(
-    "text-generation",
-    model="sarvamai/sarvam-translate",
-    torch_dtype=torch.float32,
-    device="cuda:0",
-)
+model_name = "sarvamai/sarvam-translate"
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32).to('cuda:0')
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def generate_response(tgt_lang, user_prompt):
     messages = [
         {"role": "system", "content": f"Translate the following sentence into {tgt_lang}."},
         {"role": "user", "content": user_prompt},
     ]
     
-    output = pipe(messages, max_new_tokens=2048)
-    return output[0]["generated_text"][-1]["content"]
+    # Apply chat template to structure the conversation
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+    )
+    
+    # Tokenize and move input to model device
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    
+    # Generate the output
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=1024,
+        do_sample=True,
+        temperature=0.01,
+        num_return_sequences=1
+    )
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+    return tokenizer.decode(output_ids, skip_special_tokens=True)
 
 # Create Gradio UI
 demo = gr.Interface(