Update app.py

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from huggingface_hub.inference._generated.types import TextGenerationStreamOutput, TextGenerationStreamOutputToken
 import os
 from threading import Thread
 
@@ -37,7 +38,7 @@ tokenizer = AutoTokenizer.from_pretrained(MODELS,trust_remote_code=True)
 
 
 @spaces.GPU
-def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):
+def stream_chat(message: str, history: list, temperature: float, max_length: int):
     print(f'message is - {message}')
     print(f'history is - {history}')
     conversation = []
@@ -48,22 +49,31 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
     print(f"Conversation is -\n{conversation}")
     
     input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
-    
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
-        max_length=2500, 
-        max_new_tokens=max_new_tokens,
+        max_length=max_length,
+        streamer=streamer,
         do_sample=True,
         top_k=1,
         temperature=temperature,
         repetition_penalty=1.2,
     )
+    gen_kwargs = {**input_ids, **generate_kwargs}
 
     with torch.no_grad():
-        outputs = model.generate(**inputs, **generate_kwargs)
-        outputs = outputs[:, inputs['input_ids'].shape[1]:]
-        results = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return results
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+        for next_text in streamer:
+            yield TextGenerationStreamOutput(
+                index=0,
+                token=TextGenerationStreamOutputToken(
+                    id=0,
+                    logprob=0,
+                    text=next_text,
+                    special=False,
+                )
+            )
         
 
 
@@ -90,10 +100,10 @@ with gr.Blocks(css=CSS) as demo:
             ),
             gr.Slider(
                 minimum=128,
-                maximum=4096,
+                maximum=8192,
                 step=1,
                 value=1024,
-                label="Max new tokens",
+                label="Max Length",
                 render=False,
             ),
         ],