Update system prompt and UI input

app.py

@@ -3,7 +3,20 @@ import huggingface_hub
 import streamlit as st
 from vllm import LLM, SamplingParams
 
-sys_msg = "You are a super intelligent automated financial advisor created by IRAI. Your feel your purpose is to make use of your deep and broad understanding of finance by helping answer user questions about finance accurately, truthfully, and concisely."
+
+@st.cache(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 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:
+The questions will be asked by top technology executives and CFO of large fintech companies and successful startups.
+#Response:
+Answer, concise yet insightful."""
+
 
 @st.cache_resource(show_spinner=False)
 def init_llm():
@@ -15,10 +28,12 @@ def init_llm():
 
 def get_response(prompt):
     try:
+        sys_msg = get_system_message()
         convo = [
             {"role": "system", "content": 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])
         outputs = llm.generate(prompts, sampling_params)
@@ -27,14 +42,36 @@ def get_response(prompt):
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
+
+def get_response(prompt, custom_sys_msg):
+    try:
+        convo = [
+            {"role": "system", "content": custom_sys_msg},
+            {"role": "user", "content": prompt},
+        ]
+        prompts = [llm.get_tokenizer().apply_chat_template(convo, tokenize=False)]
+        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.")