Update to use latest model

app.py

@@ -1,3 +1,4 @@
+import torch
 import streamlit as st
 from transformers import pipeline
 from concurrent.futures import ThreadPoolExecutor
@@ -6,42 +7,52 @@ from concurrent.futures import ThreadPoolExecutor
 # Function to load models only once using Streamlit's cache mechanism
 @st.cache_resource(show_spinner="Loading Models...")
 def load_models():
+    device = 0 if torch.cuda.is_available() else -1
     base_pipe = pipeline(
         "text-generation",
         model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-        max_length=512,
+        device=device,
     )
     irai_pipe = pipeline(
         "text-generation",
         model="InvestmentResearchAI/LLM-ADE_tiny-v0.001",
-        max_length=512,
+        device=device,
     )
     return base_pipe, irai_pipe
 
 
 base_pipe, irai_pipe = load_models()
 
-prompt_template = (
+alpaca_template = (
     "<|system|>\n"
-    "You are a friendly chatbot who always gives helpful, detailed, and polite answers.</s>\n"
+    "{sys}</s>\n"
     "<|user|>\n"
     "{input_text}</s>\n"
     "<|assistant|>\n"
 )
 
+chatml_template = (
+    "<|im_start|>system\n"
+    "{sys}<|im_end|>\n"
+    "<|im_start|>user\n"
+    "{input_text}<|im_end|>\n"
+    "<|im_start|>assistant\n"
+)
+
+system_prompt = "You are an AI assistant trained on an extensive dataset, including technology reports, investment reports, financial texts, and other relevant sources. Please answer the following question based on the knowledge you have acquired during your training. Do not make any assumptions or use information from external sources. If you don't have enough pre-existing knowledge to provide a complete answer, simply respond with \"I don't have enough pre-existing knowledge to comprehensively answer this question.\" If you can partially answer the question based on your training, please provide that partial answer and clarify that it may not be a complete response. Assume today is June 5, 2024, and respond as if you have no knowledge of events after your training data's cut-off date."
 executor = ThreadPoolExecutor(max_workers=2)
 
 
 def generate_base_response(input_text):
-    formatted_input = prompt_template.format(input_text=input_text)
+    formatted_input = alpaca_template.format(sys=system_prompt, input_text=input_text)
     result = base_pipe(formatted_input)[0]["generated_text"]
     return result.split("<|assistant|>")[1].strip()
 
 
 def generate_irai_response(input_text):
-    formatted_input = prompt_template.format(input_text=input_text)
+    formatted_input = chatml_template.format(sys=system_prompt, input_text=input_text)
     result = irai_pipe(formatted_input)[0]["generated_text"]
-    return result.split("<|assistant|>")[1].strip()
+    return result.split("<|im_start|>assistant")[1].split("<|im_end|>")[0].strip()
 
 
 @st.cache_data(show_spinner="Generating responses...")
@@ -49,7 +60,7 @@ def generate_response(input_text):
     try:
         future_base = executor.submit(generate_base_response, input_text)
         future_irai = executor.submit(generate_irai_response, input_text)
-        base_resp = future_base.result().replace(input_text, "", 1).strip()
+        base_resp = future_base.result()
         irai_resp = future_irai.result()
     except Exception as e:
         st.error(f"An error occurred: {e}")
@@ -58,17 +69,18 @@ def generate_response(input_text):
 
 
 st.title("Base Model vs IRAI LLM-ADE")
-user_input = st.text_area("Enter a financial question:", "")
+st.markdown("This is a demonstration of the [LLM-ADE paper](https://arxiv.org/abs/2404.13028) (knowledge cutoff is June 5, 2024)")
+user_input = st.text_area("Ask about finance related questions and mega-cap (top 15) stocks!", "")
 
 if st.button("Generate") or user_input:
     if user_input:
         base_response, irai_response = generate_response(user_input)
         col1, col2 = st.columns(2)
         with col1:
-            st.header("Base Model")
-            st.text_area(label="", value=base_response, height=300)
+            st.write("### Base Model (Tiny-Llama)")
+            st.text_area(label="none", value=base_response, height=300, key="base_response", label_visibility="hidden")
         with col2:
-            st.header("LLM-ADE Enhanced")
-            st.text_area(label="", value=irai_response, height=300)
+            st.write("### LLM-ADE Enhanced")
+            st.text_area(label="none", value=irai_response, height=300, key="irai_response", label_visibility="hidden")
     else:
         st.warning("Please enter some text")