Update app.py

app.py

@@ -1,5 +1,5 @@
 from threading import Thread
-from transformers import TextStreamer
+from transformers import TextStreamer, TextIteratorStreamer
 from unsloth import FastLanguageModel
 import torch
 import gradio as gr
@@ -21,7 +21,10 @@ FastLanguageModel.for_inference(model)
 print("model loaded")
 
 
-streamer = TextStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens = True)
+streamer = TextIteratorStreamer(
+        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+    )
+# streamer = TextStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens = True)
 
 messages = []
 
@@ -36,25 +39,26 @@ def generate_text(prompt, max_length, top_p, top_k):
     )
 
 
-    # generate_kwargs = dict(
-    #     input_ids,
-    #     max_length=int(max_length),top_p=float(top_p), do_sample=True, top_k=int(top_k), streamer=streamer
-    # )
+    generate_kwargs = dict(
+        max_length=int(max_length),top_p=float(top_p), do_sample=True,
+        top_k=int(top_k), streamer=streamer, temperature=0.6, repetition_penalty=1.2
+        
+    )
 
-    _ = model.generate(input_ids, streamer = streamer, max_new_tokens = int(max_length), pad_token_id = tokenizer.eos_token_id,
-        temperature=0.6,  # Adjust this value
-        top_k=int(top_k),        # Adjust this value
-        top_p=float(top_p),       # Adjust this value
-        repetition_penalty=1.2
-                       )
-    # t = Thread(target=model.generate, kwargs=generate_kwargs)
-    # t.start()
+    # _ = model.generate(input_ids, streamer = streamer, max_new_tokens = int(max_length), pad_token_id = tokenizer.eos_token_id,
+    #     temperature=0.6,  # Adjust this value
+    #     top_k=int(top_k),        # Adjust this value
+    #     top_p=float(top_p),       # Adjust this value
+    #     repetition_penalty=1.2
+    #                    )
+    t = Thread(target=model.generate,  args=(input_ids,), kwargs=generate_kwargs)
+    t.start()
 
     generated_text=[]
 
     for text in streamer:
         generated_text.append(text)
-        print(generated_text)
+        # print(generated_text)
         yield "".join(generated_text)
     
     messages.append({"role": "assistant", "content": "".join(generated_text)})