Update app.py

app.py

@@ -24,9 +24,6 @@ translator = pipeline('translation', model=trans_model, tokenizer=eng_trans_toke
 # Initialize translation pipelines
 pipe = pipeline("translation", model="thilina/mt5-sinhalese-english")
 
-trans_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
-eng_trans_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")
-
 sin_trans_model = AutoModelForSeq2SeqLM.from_pretrained("thilina/mt5-sinhalese-english")
 si_trans_tokenizer = AutoTokenizer.from_pretrained("thilina/mt5-sinhalese-english")
 
@@ -102,12 +99,13 @@ def transliterate_to_sinhala(text):
   latin_text = transliterate.process(source_script, target_script, text)
   return latin_text
 
-ai_model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2-0.5B-Instruct-GPTQ-Int4",
-    torch_dtype="auto",
-    device_map="auto"
+model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Phi-3-mini-4k-instruct", 
+    device_map="cuda", 
+    torch_dtype="auto", 
+    trust_remote_code=True, 
 )
-ai_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct-GPTQ-Int4")
+tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
 
     
 def conversation_predict(prompt):
@@ -115,24 +113,22 @@ def conversation_predict(prompt):
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": prompt}
     ]
-    text = ai_tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
     )
-    model_inputs = ai_tokenizer([text], return_tensors="pt").to(device)
     
-    generated_ids = ai_model.generate(
-        model_inputs.input_ids,
-        max_new_tokens=512
-    )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
+    generation_args = {
+        "max_new_tokens": 500,
+        "return_full_text": False,
+        "temperature": 0.0,
+        "do_sample": False,
+    }
     
-    response = ai_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    output = pipe(messages, **generation_args)
 
-    return response
+    return output[0]['generated_text']
 
 def ai_predicted(user_input):
     user_input = translate_Singlish_to_sinhala(user_input)