"Update space"

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

app.py

@@ -102,16 +102,17 @@ 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 with streaming support
 def mistral_response_stream(user_input):
     try:
         inputs = tokenizer(user_input, return_tensors="pt", truncation=True)
+        outputs = model.generate(inputs["input_ids"], max_length=150, do_sample=True, temperature=0.7,
+                                 pad_token_id=tokenizer.eos_token_id)
+
         # 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)
+        for output in outputs:
+            response = tokenizer.decode(output, skip_special_tokens=True)
+            yield response
     except Exception as e:
         yield f"Error generating response: {str(e)}"
 
@@ -167,10 +168,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="green", neutral_hue="slate")) a
     submit_button.click(
         fn=chatbot_response,
         inputs=[fecal_image, user_input],
-        outputs=[output_box],
-        stream=True  # Enable streaming
+        outputs=[output_box]
     )
 
 # Launch the Gradio interface
 if __name__ == "__main__":
-    chatbot_interface.launch(debug=True, share=True)
+    chatbot_interface.queue().launch(debug=True, share=True)