Update app.py

app.py

@@ -241,7 +241,6 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
         conversation = [{"role": "user", "content": input_text}]
         input_ids = tokenizer(conversation[-1]['content'], return_tensors="pt").to(model.device)
-        streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
         generate_kwargs = dict(
             max_length=max_length,
@@ -250,30 +249,16 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             top_k=top_k,
             temperature=temperature,
             repetition_penalty=penalty,
-            eos_token_id=[151329, 151336, 151338],
-            streamer=streamer
+            eos_token_id=[151329, 151336, 151338]
         )
 
-        buffer = ""
-
-        def generate_text():
-            with torch.no_grad():
-                model.generate(input_ids['input_ids'], **generate_kwargs)
-
-        # Run generation in a separate thread
-        thread = Thread(target=generate_text)
-        thread.start()
-
-        # Collect generated text in real-time
-        for new_text in streamer:
-            buffer += new_text
-
-        # Wait for the generation thread to finish
-        thread.join()
+        with torch.no_grad():
+            generated_ids = model.generate(input_ids['input_ids'], **generate_kwargs)
+            generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
 
         # Process to remove any prefix or unwanted prompt
         text_original = input_text.strip()
-        results_text = buffer[len(text_original):].strip()
+        results_text = generated_text[len(text_original):].strip()
 
         print("  ")
         print("------")
@@ -289,7 +274,6 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
 
 
-
 # @spaces.GPU()
 # def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
 #     try: