Reopen preprompt UI

app.py

@@ -3,10 +3,13 @@ import huggingface_hub
 import streamlit as st
 from vllm import LLM, SamplingParams
 
-sys_msg = """#Context:
-You are an expert financial advisor named IRAI. You have a comprehensive understanding of finance and investing with experience and expertise in all areas of finance. 
+
+@st.cache_data(show_spinner=False)
+def get_system_message():
+    return """#Context:
+You are an AI-based automated expert financial advisor named IRAI. You have a comprehensive understanding of finance and investing because you have trained on a  extensive dataset based on of financial news, analyst reports, books, company filings, earnings call transcripts, and finance websites.
 #Objective:
-Answer questions accurately and truthfully given your current knowledge.  You do not have access to up-to-date current market data; this will be available in the future. 
+Answer questions accurately and truthfully given the data you have trained on.  You do not have access to up-to-date current market data; this will be available in the future. 
 Style and tone:
 Please answer in a friendly and engaging manner representing a top female investment professional working at a leading investment bank.
 #Audience:
@@ -14,7 +17,8 @@ The questions will be asked by top technology executives and CFO of large fintec
 #Response:
 Answer, concise yet insightful."""
 
-@st.cache_resource(show_spinner="Loading model..")
+
+@st.cache_resource(show_spinner=False)
 def init_llm():
     huggingface_hub.login(token=os.getenv("HF_TOKEN"))
     llm = LLM(model="InvestmentResearchAI/LLM-ADE-dev")
@@ -22,31 +26,36 @@ def init_llm():
     tok.eos_token = '<|im_end|>' # Override to use turns
     return llm
 
-def get_response(prompt):
+
+def get_response(prompt, custom_sys_msg):
     try:
         convo = [
-            {"role": "system", "content": sys_msg},
+            {"role": "system", "content": custom_sys_msg},
             {"role": "user", "content": prompt},
         ]
-        llm = init_llm()
         prompts = [llm.get_tokenizer().apply_chat_template(convo, tokenize=False)]
-        sampling_params = SamplingParams(temperature=0.3, top_p=0.95, max_tokens=500, stop_token_ids=[128009])
+        sampling_params = SamplingParams(temperature=0.3, top_p=0.95, max_tokens=2000, stop_token_ids=[128009])
         outputs = llm.generate(prompts, sampling_params)
         for output in outputs:
             return output.outputs[0].text
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-
 def main():
     st.title("LLM-ADE 9B Demo")
     
+    # Retrieve the default system message
+    sys_msg = get_system_message()
+    
+    # UI for editable preprompt
+    user_modified_sys_msg = st.text_area("Preprompt: ", value=sys_msg, height=200)
+
     input_text = st.text_area("Enter your text here:", value="", height=200)
     
     if st.button("Generate"):
         if input_text:
             with st.spinner('Generating response...'):
-                response_text = get_response(input_text)
+                response_text = get_response(input_text, user_modified_sys_msg)
                 st.write(response_text)
         else:
             st.warning("Please enter some text to generate a response.")