Update app.py

app.py

@@ -7,7 +7,7 @@ from ollama import ChatResponse
 # Default model
 OLLAMA_MODEL = "llama3.2:3b"
 # OLLAMA_MODEL = "llama3.2:1b"
-OLLAMA_MODEL = "llama3.2:3b-instruct-q2_K"
+# OLLAMA_MODEL = "llama3.2:3b-instruct-q2_K"
 
 # Load BERT MODEL
 from transformers import pipeline, DistilBertTokenizerFast
@@ -29,10 +29,15 @@ bert_label_map = {
     'LABEL_1': 'NTA',
 }
 
+bert_label_map_formatted = {
+    'LABEL_0': 'You are the A**hole (YTA)',
+    'LABEL_1': 'Not the A**hole (NTA)',
+}
+
 def ask_bert(prompt):
     print(f"Getting response from Fine-tuned BERT")
     result = classifier([prompt])[0]
-    label = bert_label_map.get(result['label'])
+    label = result['label']
     confidence = f"{result['score']*100:.2f}"
     return label, confidence
 
@@ -99,9 +104,6 @@ Use second person terms like you in the explanation.
         response += chunk['message']['content']
         yield response
 
-def gradio_bert_interface(prompt):
-    response, confidence = ask_bert(prompt)
-    return f"{response} with confidence {confidence}%"
 
 # Separate function for Ollama response
 def gradio_ollama_interface(prompt, bert_class=""):
@@ -111,14 +113,15 @@ def gradio_interface(prompt, selected_model):
         for chunk in ask_ollama(prompt):
             yield chunk
     elif selected_model == MODEL_CHOICE_BERT:
-        response, confidence = ask_bert(prompt)
-        response = f"{response} with confidence {confidence}"
+        label, confidence = ask_bert(prompt)
+        label = bert_label_map_formatted[label]
+        response = f"{label} with confidence {confidence}"
         return response
     elif selected_model == MODEL_CHOICE_BERT_LLAMA:
         label, confidence = ask_bert(prompt)
-        initial_response = f"BERT model says {label} with confidence {confidence}%\n\nGenerating explanation using Llama model...\n"
+        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_ollama(prompt, expected_class=label):
+        for chunk in ask_ollama(prompt, expected_class=bert_label_map[label]):
             yield initial_response + "\n" + chunk
     else:
         return "Something went wrong. Select the correct model configuration from settings. "