Update app.py

app.py

@@ -102,23 +102,37 @@ def transliterate_to_sinhala(text):
   latin_text = transliterate.process(source_script, target_script, text)
   return latin_text
 
-# Placeholder for conversation model loading and pipeline setup
-# pipe1 = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)
-
-# interface = gr.Interface.load("huggingface/microsoft/Phi-3-mini-4k-instruct")
-
-# API_URL = "https://api-inference.huggingface.co/models/microsoft/Phi-3-mini-4k-instruct"
-# headers = {"Authorization": f"Bearer {access_token}"}
-
-ai_pipe = pipeline("text-generation", model="google/gemma-2b-it", token = access_token)
+ai_model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2-0.5B-Instruct-GPTQ-Int4",
+    torch_dtype="auto",
+    device_map="auto"
+)
+ai_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct-GPTQ-Int4")
 
-# def query(payload):
-# 	response = requests.post(API_URL, headers=headers, json=payload)
-# 	return response.json()
     
-def conversation_predict(text):
-    ai_response = ai_pipe([{"role": "user", "content": f"{text}"}])
-    return ai_response
+def conversation_predict(prompt):
+    messages = [
+        {"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
+    )
+    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)
+    ]
+    
+    response = ai_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    return response
 
 def ai_predicted(user_input):
     user_input = translate_Singlish_to_sinhala(user_input)