Update app.py

app.py

@@ -44,14 +44,28 @@ def transliterate_to_sinhala(text):
     return transliterate.process('Velthuis', 'Sinhala', text)
 
 # Load conversation model
-conv_model_name = "Qwen/Qwen2-0.5B-Instruct"  # Use GPT-2 instead of the gated model
-tokenizer = AutoTokenizer.from_pretrained(conv_model_name)
+conv_model_name = "microsoft/Phi-3-mini-4k-instruct"  # Use GPT-2 instead of the gated model
+tokenizer = AutoTokenizer.from_pretrained(conv_model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(conv_model_name).to(device)
 
 def conversation_predict(text):
-    input_ids = tokenizer(text, return_tensors="pt").to(device)
-    outputs = model.generate(**input_ids)
-    return tokenizer.decode(outputs[0])
+    pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    )
+    generation_args = {
+    "max_new_tokens": 500,
+    "return_full_text": False,
+    "temperature": 0.0,
+    "do_sample": False,
+    }
+
+    output = pipe(text, **generation_args)
+    return output[0]['generated_text']
+    # input_ids = tokenizer(text, return_tensors="pt").to(device)
+    # outputs = model.generate(**input_ids)
+    # return tokenizer.decode(outputs[0])
 
 def ai_predicted(user_input):
     if user_input.lower() == 'exit':