Update app.py

app.py

@@ -1,13 +1,15 @@
 import spaces
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 import torch
 
-model_name = "sarvamai/sarvam-translate"
-
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to('cuda:0')
+# Initialize the pipeline
+pipe = pipeline(
+    "text-generation",
+    model="sarvamai/sarvam-translate",
+    torch_dtype=torch.float32,
+    device="cuda:0",
+)
 
 @spaces.GPU
 def generate_response(tgt_lang, user_prompt):
@@ -16,25 +18,8 @@ def generate_response(tgt_lang, user_prompt):
         {"role": "user", "content": user_prompt},
     ]
     
-    # 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)
+    output = pipe(messages, max_new_tokens=2048)
+    return output[0]["generated_text"][-1]["content"]
 
 # Create Gradio UI
 demo = gr.Interface(