Update app.py

app.py

@@ -3,6 +3,15 @@ import subprocess
 import time
 from ollama import chat
 from ollama import ChatResponse
+from huggingface_hub import InferenceClient
+import os
+
+# HF client
+hf_api_key = os.getenv("HF_API_KEY", None)
+if not hf_api_key:
+    print("HF_API_KEY environment variable not set, will not be able to use Inference API")
+hf_client = InferenceClient(api_key=hf_api_key)
+
 
 # Default model
 OLLAMA_MODEL = "llama3.2:3b"
@@ -105,6 +114,65 @@ Use second person terms like you in the explanation.
         yield response
 
 
+
+def ask_hf_inference_client(question, expected_class=""):
+    print(f"Getting response from Ollama")
+    classify_and_explain_prompt = f"""
+### You are an unbiased expert from subreddit community r/AmItheAsshole. In this community people post their life situations and ask if they are the asshole or not. 
+The community uses the following acronyms. 
+AITA : Am I the asshole? Usually posted in the question. 
+YTA : You are the asshole in this situation.  
+NTA : You are not the asshole in this situation.
+
+### The task for you label YTA or NTA for the given text. Give a short explanation for the label. Be brutally honest and unbiased. Base your explanation entirely on the given text only. 
+
+If the label is YTA, also explain what could the user have done better.  
+### The output format is as follows:
+"YTA" or "NTA", a short explanation. 
+
+### Situation :  {question}
+### Response :"""
+
+    explain_only_prompt =  f"""
+### You know about the subreddit community r/AmItheAsshole. In this community people post their life situations and ask if they are the asshole or not. 
+The community uses the following acronyms. 
+AITA : Am I the asshole? Usually posted in the question. 
+YTA : You are the asshole in this situation.  
+NTA : You are not the asshole in this situation.
+
+### The task for you explain why a particular situation was tagged as NTA or YTA by most users. I will give the situation as well as the NTA or YTA tag. just give your explanation for the label. Be nice but give a brutally honest and unbiased view. Base your explanation entirely on the given text and the label tag only. Do not assume anything extra.  
+Use second person terms like you in the explanation.
+
+### Situation :  {question}
+### Label Tag : {expected_class}
+### Explanation for {expected_class} :"""
+
+    if expected_class == "":
+        prompt = classify_and_explain_prompt
+    else:
+        prompt = explain_only_prompt
+
+    print(f"Prompt to HF_Inference API : {prompt}")
+
+    messages = [
+        {
+            "role": "user",
+            "content": prompt
+        }
+    ]
+
+    stream = hf_client.chat.completions.create(
+        model="meta-llama/Llama-3.2-3B-Instruct",
+        messages=messages,
+        max_tokens=500,
+        stream=True
+    )
+
+    for chunk in stream:
+        yield chunk.choices[0].delta.content
+
+
+
 # Separate function for Ollama response
 def gradio_ollama_interface(prompt, bert_class=""):
     return ask_ollama(prompt, expected_class=bert_class)
@@ -123,14 +191,22 @@ def gradio_interface(prompt, selected_model):
         yield initial_response
         for chunk in ask_ollama(prompt, expected_class=bert_label_map[label]):
             yield initial_response + "\n" + chunk
+    elif selected_model == MODEL_CHOICE_BERT_LLAMA_HF_INFERENCE:
+        label, confidence = ask_bert(prompt)
+        initial_response = f"Response from BERT model:  {bert_label_map_formatted[label]} with confidence {confidence}%\n\nGenerating explanation using Llama model...\n"
+        yield initial_response
+        for chunk in ask_hf_inference_client(prompt, expected_class=bert_label_map[label]):
+            initial_response += chunk
+            yield initial_response
     else:
         return "Something went wrong. Select the correct model configuration from settings. "
 
 MODEL_CHOICE_BERT_LLAMA = "Fine-tuned BERT (classification) + Llama 3.2 3B (explanation)"
+MODEL_CHOICE_BERT_LLAMA_HF_INFERENCE = "Fine-tuned BERT (classification) + Llama 3.2 3B Inference api (fast explanation)"
 MODEL_CHOICE_BERT = "Fine-tuned BERT (classification only)"
 MODEL_CHOICE_LLAMA = "Llama 3.2 3B (classification + explanation)"
 
-MODEL_OPTIONS = [MODEL_CHOICE_BERT_LLAMA, MODEL_CHOICE_LLAMA, MODEL_CHOICE_BERT]
+MODEL_OPTIONS = [MODEL_CHOICE_BERT_LLAMA_HF_INFERENCE, MODEL_CHOICE_BERT_LLAMA, MODEL_CHOICE_LLAMA, MODEL_CHOICE_BERT]
 
 # Example texts
 EXAMPLES = [
@@ -186,5 +262,5 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.green, secon
 
 # Launch the app
 if __name__ == "__main__":
-    start_ollama_server()
+    # start_ollama_server()
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)