WilliamGazeley commited on
Commit
e87746b
·
1 Parent(s): a20dfac

Add streamlit resource cache

Browse files
Files changed (1) hide show
  1. app.py +9 -5
app.py CHANGED
@@ -3,11 +3,6 @@ import huggingface_hub
3
  import streamlit as st
4
  from vllm import LLM, SamplingParams
5
 
6
- huggingface_hub.login(token=os.getenv("HF_TOKEN"))
7
- llm = LLM(model="InvestmentResearchAI/LLM-ADE-small-v0.1.0")
8
- tok = llm.get_tokenizer()
9
- tok.eos_token = '<|eot_id|>' # Override to use turns
10
-
11
 
12
  template = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
13
 
@@ -17,6 +12,13 @@ You are a helpful financial assistant that answers the user as accurately, truth
17
 
18
  """
19
 
 
 
 
 
 
 
 
20
 
21
  def get_response(prompt):
22
  try:
@@ -40,6 +42,8 @@ def main():
40
  else:
41
  st.warning("Please enter some text to generate a response.")
42
 
 
 
43
  if __name__ == "__main__":
44
  main()
45
 
 
3
  import streamlit as st
4
  from vllm import LLM, SamplingParams
5
 
 
 
 
 
 
6
 
7
  template = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
8
 
 
12
 
13
  """
14
 
15
+ @st.cache_resource
16
+ def init_llm():
17
+ huggingface_hub.login(token=os.getenv("HF_TOKEN"))
18
+ llm = LLM(model="InvestmentResearchAI/LLM-ADE-small-v0.1.0")
19
+ tok = llm.get_tokenizer()
20
+ tok.eos_token = '<|eot_id|>' # Override to use turns
21
+ return llm
22
 
23
  def get_response(prompt):
24
  try:
 
42
  else:
43
  st.warning("Please enter some text to generate a response.")
44
 
45
+ llm = init_llm()
46
+
47
  if __name__ == "__main__":
48
  main()
49