Implement simple comparison UI

app.py

@@ -1,24 +1,47 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
+from concurrent.futures import ProcessPoolExecutor
 
-# Load the model and tokenizer
-model_name = "InvestmentResearchAI/LLM-ADE_tiny-v0.001"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+prompt_template = (
+    "<|system|>\n"
+    "You are a friendly chatbot who always gives helpful, detailed, and polite answers.</s>\n"
+    "<|user|>\n"
+    "{input_text}</s>\n"
+    "<|assistant|>\n"
+)
+
+def generate_base_response(input_text):
+    base_pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", max_length=512)
+    return base_pipe(input_text)[0]["generated_text"]
+
+def generate_irai_response(input_text):
+    irai_pipe = pipeline("text-generation", model="InvestmentResearchAI/LLM-ADE_tiny-v0.001", max_length=512)
+    return irai_pipe(prompt_template.format(input_text=input_text))[0]["generated_text"].split("<|assistant|>")[1].strip()
 
 def generate_response(input_text):
-    """Generate response from the model based on the input text."""
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    output = model.generate(**inputs, max_length=512, num_return_sequences=1)
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
+    with ProcessPoolExecutor() as executor:
+        try:
+            future_base = executor.submit(generate_base_response, input_text)
+            future_irai = executor.submit(generate_irai_response, input_text)
+            base_resp = future_base.result()
+            irai_resp = future_irai.result()
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+            return None, None
+    return base_resp, irai_resp
 
-# Streamlit interface
-st.title("IRAI LLM-ADE Model")
-user_input = st.text_area("Enter your text here:", "")
+st.title("IRAI LLM-ADE Model vs Base Model")
+user_input = st.text_area("Enter a financial question:", "")
 if st.button("Generate"):
     if user_input:
-        response = generate_response(user_input)
-        st.text_area("Model Response:", response, height=300)
+        base_response, irai_response = generate_response(user_input)
+        col1, col2 = st.columns(2)  # Updated to use `st.columns`
+        with col1:
+            st.header("Base Model Response")
+            st.text_area("", base_response, height=300)
+        with col2:
+            st.header("IRAI LLM-ADE Model Response")
+            st.text_area("", irai_response, height=300)
     else:
-        st.warning("Please enter some text to generate a response.")
+        st.warning("Please enter some text to generate a response.")
+