async support

app.py

@@ -5,9 +5,8 @@ from ollama import chat
 from ollama import ChatResponse
 
 # Default model
-# OLLAMA_MODEL = "llama3.2:3b-instruct-q3_K_M"
-# OLLAMA_MODEL = "llama3.2:1b"
 OLLAMA_MODEL = "llama3.2:3b"
+# OLLAMA_MODEL = "llama3.2:1b"
 
 # Load BERT MODEL
 from transformers import pipeline, DistilBertTokenizerFast
@@ -88,28 +87,40 @@ Use second person terms like you in the explanation.
         prompt = explain_only_prompt
 
     print(f"Prompt to llama : {prompt}")
-    response: ChatResponse = chat(model=OLLAMA_MODEL, messages=[
+    stream = chat(model=OLLAMA_MODEL, messages=[
         {
             'role': 'user',
             'content': prompt,
         },
-    ])
-    print(response['message']['content'])
-    return response['message']['content']
-
+    ], stream=True)
+    response = ""
+    for chunk in stream:
+        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=""):
+    return ask_ollama(prompt, expected_class=bert_class)
 def gradio_interface(prompt, selected_model):
     if selected_model == MODEL_CHOICE_LLAMA:
-        response = ask_ollama(prompt)
+        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}"
+        return response
     elif selected_model == MODEL_CHOICE_BERT_LLAMA:
-        bert_response, confidence = ask_bert(prompt)
-        ollama_response = ask_ollama(prompt, expected_class=bert_response)
-        response = f"{bert_response} with {confidence}% confidence. \n {ollama_response}"
+        label, confidence = ask_bert(prompt)
+        initial_response = f"BERT model says {label} with confidence {confidence}%\n\nGenerating explanation using Llama model...\n"
+        yield initial_response
+        for chunk in ask_ollama(prompt, expected_class=label):
+            yield initial_response + "\n" + chunk
     else:
-        response = "Something went wrong. Select the correct model configuration from settings. "
-    return response
+        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 = "Fine-tuned BERT (classification only)"