changes on app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from peft import PeftModel
 import re
 import os
@@ -12,6 +12,7 @@ import math
 from datetime import date
 from datetime import date, datetime, timedelta
 
+
 access_token = os.environ["TOKEN"]
 
 # load model
@@ -22,7 +23,9 @@ tokenizer = AutoTokenizer.from_pretrained(model, token = access_token, trust_rem
 tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = "right"
 
-model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True, token = access_token, device_map="auto", load_in_8bit=True, offload_folder="offload/")
+streamer = TextStreamer(tokenizer)
+
+model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True, token = access_token, device_map="cuda", load_in_8bit=True, offload_folder="offload/")
 model = PeftModel.from_pretrained(model, peft_model, offload_folder="offload/")
 
 model = model.eval()
@@ -366,7 +369,10 @@ def get_all_prompts_online(symbol, with_basics=True, max_news_perweek = 3, weeks
     new_system_prompt = SYSTEM_PROMPT.replace('：\n...', '：\n预测涨跌幅：...\n总结分析：...')
     prompt = B_INST + B_SYS + new_system_prompt + E_SYS + info + f"\n\n基于在{end_date}之前的所有信息，让我们首先分析{stock}的积极发展和潜在担忧。请简洁地陈述，分别提出2-4个最重要的因素。大部分所提及的因素应该从公司的相关新闻中推断出来。" \
         f"接下来请预测{symbol}下周({period})的股票涨跌幅，并提供一个总结分析来支持你的预测。" + E_INST
-        
+    
+    del prev_rows
+    del data
+
     return info, prompt
 
 
@@ -382,7 +388,9 @@ def ask(symbol, weeks_before):
 
   res = model.generate(
       **inputs,
-      use_cache=True
+      use_cache=True,
+      max_length = 4096,
+      streamer=streamer
   )
   output = tokenizer.decode(res[0], skip_special_tokens=True)
   output_cur = re.sub(r'.*\[/INST\]\s*', '', output, flags=re.DOTALL)