Update app.py

app.py

@@ -4,12 +4,14 @@ import time
 import json
 import random
 import finnhub
+import torch
 import gradio as gr
 import pandas as pd
 import yfinance as yf
-from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+from collections import defaultdict
 from datetime import date, datetime, timedelta
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 os.environ['HF_HOME'] = '/data/.huggingface'
 
@@ -36,7 +38,6 @@ tokenizer = AutoTokenizer.from_pretrained(
     token=access_token
 )
 
-
 B_INST, E_INST = "[INST]", "[/INST]"
 B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
 
@@ -208,17 +209,18 @@ def get_all_prompts_online(symbol, data, curday, with_basics=True):
     return info, prompt
 
 
-def construct_prompt(ticker, date, n_weeks, use_basics):
+def construct_prompt(ticker, curday, n_weeks, use_basics):
 
-    curday = get_curday()
     steps = [n_weeks_before(curday, n) for n in range(n_weeks + 1)][::-1]
     data = get_stock_data(ticker, steps)
     data = get_news(ticker, data)
     data['Basics'] = [json.dumps({})] * len(data)
+    print(data)
     
     info, prompt = get_all_prompts_online(ticker, data, curday, use_basics)
     
     prompt = B_INST + B_SYS + SYSTEM_PROMPT + E_SYS + prompt + E_INST
+    print(prompt)
     
     return info, prompt
 
@@ -228,8 +230,7 @@ def predict(ticker, date, n_weeks, use_basics):
     info, prompt = construct_prompt(ticker, date, n_weeks, use_basics)
       
     inputs = tokenizer(
-        prompt, return_tensors='pt',
-        padding=False, max_length=4096
+        prompt, return_tensors='pt', padding=False
     )
     inputs = {key: value.to(model.device) for key, value in inputs.items()}
         
@@ -240,6 +241,8 @@ def predict(ticker, date, n_weeks, use_basics):
     )
     output = tokenizer.decode(res[0], skip_special_tokens=True)
     answer = re.sub(r'.*\[/INST\]\s*', '', output, flags=re.DOTALL)
+
+    torch.cuda.empty_cache()
     
     return info, answer