"Update space"

Signed-off-by: Emmanuel Frimpong Asante <[email protected]>

app.py

@@ -102,22 +102,18 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Define Mistral-based response generation
-def mistral_response(user_input):
+
+# Define Mistral-based response generation with streaming support
+def mistral_response_stream(user_input):
     try:
-        responses = generator(
-            user_input,
-            max_length=150,
-            num_return_sequences=1,
-            truncation=True,  # Explicitly set truncation
-            temperature=0.7,
-            do_sample=True,  # Enable sampling for temperature setting
-            pad_token_id=tokenizer.eos_token_id  # Set padding token ID to EOS token
-        )
-        return responses[0]["generated_text"]
+        inputs = tokenizer(user_input, return_tensors="pt", truncation=True)
+        # Stream the response token by token
+        for output in model.generate(inputs["input_ids"], max_length=150, do_sample=True, temperature=0.7,
+                                     pad_token_id=tokenizer.eos_token_id, early_stopping=True,
+                                     return_dict_in_generate=True, output_scores=True):
+            yield tokenizer.decode(output.tolist(), skip_special_tokens=True)
     except Exception as e:
-        return f"Error generating response: {str(e)}"
-
+        yield f"Error generating response: {str(e)}"
 
 # Main chatbot function: handles both generative AI and disease detection
 def chatbot_response(image, text):
@@ -125,13 +121,12 @@ def chatbot_response(image, text):
     if image is not None:
         diagnosis, name, status, recom = bot.diagnose_disease(image)
         if name and status and recom:
-            return diagnosis
+            yield diagnosis
         else:
-            return diagnosis  # Return only the diagnostic message if no disease found
-
-    # If no image is provided, proceed with generative AI response
-    return mistral_response(text)
-
+            yield diagnosis  # Return only the diagnostic message if no disease found
+    else:
+        # Stream the generative AI response
+        yield from mistral_response_stream(text)
 
 # Gradio interface styling and layout with ChatGPT-like theme
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="green", neutral_hue="slate")) as chatbot_interface:
@@ -173,6 +168,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="green", neutral_hue="slate")) a
         fn=chatbot_response,
         inputs=[fecal_image, user_input],
         outputs=[output_box],
+        stream=True  # Enable streaming
     )
 
 # Launch the Gradio interface