Update app.py

app.py

@@ -99,36 +99,19 @@ def transliterate_to_sinhala(text):
   latin_text = transliterate.process(source_script, target_script, text)
   return latin_text
 
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
 model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Phi-3-mini-4k-instruct", 
-    device_map="cuda", 
-    torch_dtype="auto", 
-    trust_remote_code=True, 
+    "google/gemma-2b-it",
+    torch_dtype=torch.bfloat16
 )
-tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
 
     
-def conversation_predict(prompt):
-    messages = [
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": prompt}
-    ]
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-    )
+def conversation_predict(input_text):
+    input_ids = tokenizer(input_text, return_tensors="pt")
     
-    generation_args = {
-        "max_new_tokens": 500,
-        "return_full_text": False,
-        "temperature": 0.0,
-        "do_sample": False,
-    }
-    
-    output = pipe(messages, **generation_args)
+    outputs = model.generate(**input_ids)
 
-    return output[0]['generated_text']
+    return tokenizer.decode(outputs[0])
 
 def ai_predicted(user_input):
     user_input = translate_Singlish_to_sinhala(user_input)