Update app.py

app.py

@@ -44,25 +44,28 @@ def transliterate_to_sinhala(text):
     return transliterate.process('Velthuis', 'Sinhala', text)
 
 # Load conversation model
-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, trust_remote_code=True).to(device)
+# 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, trust_remote_code=True).to(device)
+
+client = InferenceClient("google/gemma-2b-it")
 
 def conversation_predict(text):
-    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']
+    return client.text_generation(text, return_full_text=False)
+    # 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])