Update app.py

app.py

@@ -43,7 +43,8 @@ def check_model_status():
 
 def chat(message, history, temperature, max_new_tokens):
     global model_pipeline
-    
+    stop_tokens = ["<|endoftext|>", "<|im_end|>","|im_end|"]
+
     # Ensure the model is loaded before proceeding
     if not check_model_status():
         yield "Model is not ready. Please try again later."
@@ -53,30 +54,34 @@ def chat(message, history, temperature, max_new_tokens):
     
     # Stream the response
     start_time = time.time()
-    generated_tokens = 0
     
     # Create a TextStreamer for token streaming
     tokenizer = model_pipeline.tokenizer
     streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    def generate():
-        stop_tokens = ["<|endoftext|>", "<|im_end|>","|im_end|"]
-        nonlocal generated_tokens
-        response = model_pipeline(
-            prompt,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            do_sample=True,
-            truncation=True,
-            pad_token_id=tokenizer.eos_token_id,
-            streamer=streamer  # Use the TextStreamer here
-        )
-        for new_token in streamer:
-            outputs.append(new_token)
-            if new_token in stop_tokens:
-                
-                break
-            yield "".join(outputs), "not implemented"
+
+
+    pipeline_kwargs = dict(
+        prompt=prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        do_sample=True,
+        truncation=True,
+        pad_token_id=tokenizer.eos_token_id,
+        streamer=streamer  # Use the TextStreamer here
+    )
+    
+    # Create and start the thread with the model_pipeline function
+    t = threading.Thread(target=lambda: model_pipeline(**pipeline_kwargs))
+    t.start()
+
+    for new_token in streamer:
+        print(new_token)
+        outputs.append(new_token)
+        if new_token in stop_tokens:
+            
+            break
+        yield "".join(outputs), "not implemented"
 def reload_model_button():
     """Reload the model manually via a button."""
     global model_loaded