Update app.py

app.py

@@ -7,12 +7,17 @@ MODEL = "bigscience/bloom-560m" #"microsoft/Phi-3-mini-128k-instruct"
 
 # Initialize the InferenceClient with the token and model.
 # Using a lightweight model like google/flan-t5-small.
+# client = InferenceClient(
+#     "google/flan-t5-small",
+#     token=HF_API_TOKEN
+# )
 client = InferenceClient(
-    "google/flan-t5-small",
-    token=HF_API_TOKEN
+    provider="nebius",
+    api_key=HF_API_TOKEN
 )
 
 
+
 # -------------------------
 # Core function: call HF API
 # -------------------------
@@ -23,13 +28,25 @@ def query_llm(tweet, mode):
     else:
         prompt += "Task: Is the above TEXT toxic? First clearly state your decision. Then, provide specific reason(s) for your decision. If there is more than one reason, provide them in a numbered list. Your reason(s) must be non-redundant and jointly sufficient to justify your decision. In other words, there should not be any internal or external information unused in your explanation."
     
-    # Use the client to generate text from the full prompt.
-    response = client.text_generation(
-        prompt=prompt,
-        max_new_tokens=250,  # Control the length of the generated output.
-        stream=False,       # Set to True for streaming responses.
+    # # Use the client to generate text from the full prompt.
+    # response = client.text_generation(
+    #     prompt=prompt,
+    #     max_new_tokens=250,  # Control the length of the generated output.
+    #     stream=False,       # Set to True for streaming responses.
+    # )
+    # return response    
+
+    completion = client.chat.completions.create(
+    model="google/gemma-2-2b-it",
+    messages=[
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ],
     )
-    return response        
+    return completion.choices[0].message['content']
+    
 
 # -------------------------
 # Preloaded tweets