Update app.py

app.py

@@ -62,7 +62,7 @@ def chat(message, history, temperature, max_new_tokens):
     def generate():
         stop_tokens = ["<|endoftext|>", "<|im_end|>","|im_end|"]
         nonlocal generated_tokens
-        for response in model_pipeline(
+        response = model_pipeline(
             prompt,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
@@ -71,20 +71,12 @@ def chat(message, history, temperature, max_new_tokens):
             pad_token_id=tokenizer.eos_token_id,
             streamer=streamer  # Use the TextStreamer here
         ):
-            bot_text = response[0]["generated_text"]
-            bot_text = bot_text.split("Assistant:")[-1].strip()
-            if "</think>" in bot_text:
-                bot_text = bot_text.split("</think>")[-1].strip()
-            
-            generated_tokens += len(bot_text.split())
-            elapsed_time = time.time() - start_time
-            tokens_per_second = generated_tokens / elapsed_time if elapsed_time > 0 else 0
-            
-            yield bot_text, f"Generating... Tokens/s: {tokens_per_second:.2f}"
-    
-    for partial_response, status in generate():
-        yield partial_response, status
-
+        for new_token in streamer:
+            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