try stream

app.py

@@ -3,7 +3,7 @@ import spaces
 
 # from huggingface_hub import InferenceClient
 from transformers import pipeline
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import os
 HF_TOKEN = os.getenv('HF_TOKEN')
 
@@ -12,12 +12,17 @@ device = "cuda"  # "cuda" or "cpu"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint,token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(checkpoint,token=HF_TOKEN)
 model.to(device)
+streamer = TextStreamer(tok)
 @spaces.GPU
 def predict(message, history):
     history.append({"role": "user", "content": message})
     input_text = tokenizer.apply_chat_template(history, tokenize=False)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)  
-    outputs = model.generate(inputs, max_new_tokens=100, temperature=0.2, top_p=0.9, do_sample=True)
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device) 
+    # Use TextStreamer for streaming response
+    streamer = TextStreamer(tokenizer)
+    outputs = model.generate(inputs, max_new_tokens=512, temperature=0.2, top_p=0.9, do_sample=True, streamer=streamer)
+    # Despite returning the usual output, the streamer will also print the generated text to stdout.
+
     decoded = tokenizer.decode(outputs[0])
     response = decoded.split("[INST]")[-1].split("[/INST]")[-1]
     return response