Update app.py

app.py

@@ -15,8 +15,7 @@ alpaca_prompt = """පහත දැක්වෙන්නේ යම් කාර
 ### ප්‍රතිචාරය:
 {}"""
 
-@spaces.GPU
-def infer(message, history, enable_history=False, max_new_tokens=512):
+def infer_stream(message, history, enable_history=False, max_new_tokens=512):
     global model, tokenizer
     
     if model is None:
@@ -73,12 +72,14 @@ def infer(message, history, enable_history=False, max_new_tokens=512):
         
         if response_started:
             yield partial_text
-        elif "### ප්‍රතිචාරය:" not in prompt:
-            # If prompt doesn't contain the marker, stream everything
-            yield partial_text
     
     thread.join()
 
+@spaces.GPU
+def infer(message, history, enable_history=False, max_new_tokens=512):
+    # Return the generator for streaming
+    return infer_stream(message, history, enable_history, max_new_tokens)
+
 # Custom CSS for styling
 custom_css = """
 #splash-screen {