Update app.py

app.py

@@ -50,18 +50,19 @@ def transliterate_to_sinhala(text):
 # model = AutoModelForCausalLM.from_pretrained(conv_model_name, trust_remote_code=True).to(device)
 # pipe1 = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0").to(device)
 
-model = "tiiuae/falcon-7b-instruct"
-
-tokenizer = AutoTokenizer.from_pretrained(model)
-text_gen_pipeline = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-    device_map="auto",
-)
+# model = "tiiuae/falcon-7b-instruct"
+
+# tokenizer = AutoTokenizer.from_pretrained(model)
+# text_gen_pipeline = pipeline(
+#     "text-generation",
+#     model=model,
+#     tokenizer=tokenizer,
+#     torch_dtype=torch.bfloat16,
+#     trust_remote_code=True,
+#     device_map="auto",
+# )
 
+pipe1 = pipeline("text-generation", model="unsloth/gemma-2b-it")
 
 # client = InferenceClient("google/gemma-2b-it")
 
@@ -88,15 +89,16 @@ def conversation_predict(text):
     # outputs = pipe1(text, max_new_tokens=256, temperature=0.7, top_k=50, top_p=0.95)
     # return outputs[0]["generated_text"]
 
-    sequences = text_gen_pipeline(
-        text,
-        max_length=200,
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    return sequences[0]['generated_text']
+    # sequences = text_gen_pipeline(
+    #     text,
+    #     max_length=200,
+    #     do_sample=True,
+    #     top_k=10,
+    #     num_return_sequences=1,
+    #     eos_token_id=tokenizer.eos_token_id,
+    # )
+    # return sequences[0]['generated_text']
+    
 
 
 def ai_predicted(user_input):
@@ -133,7 +135,8 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = ai_predicted(message)
+    # response = ai_predicted(message)
+    response = pipe({"role": "user", "content": message})
 
     yield response