Update app.py

app.py

@@ -11,7 +11,7 @@ import yfinance as yf
 from peft import PeftModel
 from collections import defaultdict
 from datetime import date, datetime, timedelta
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 
 os.environ['HF_HOME'] = '/data/.huggingface'
 
@@ -38,6 +38,8 @@ tokenizer = AutoTokenizer.from_pretrained(
     token=access_token
 )
 
+streamer = TextStreamer(tokenizer)
+
 B_INST, E_INST = "[INST]", "[/INST]"
 B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
 
@@ -237,7 +239,7 @@ def predict(ticker, date, n_weeks, use_basics):
     res = model.generate(
         **inputs, max_length=4096, do_sample=True,
         eos_token_id=tokenizer.eos_token_id,
-        use_cache=True
+        use_cache=True, streamer=streamer
     )
     output = tokenizer.decode(res[0], skip_special_tokens=True)
     answer = re.sub(r'.*\[/INST\]\s*', '', output, flags=re.DOTALL)