Update app.py

app.py

@@ -5,21 +5,21 @@ import os
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
 
-# Load Hugging Face token from secrets
+# Load token from Hugging Face Space secrets
 HF_TOKEN = os.getenv("Allie", None)
 if HF_TOKEN:
     login(HF_TOKEN)
 
-# All available models
+# === Available Models for Selection ===
 model_map = {
-    "FinGPT": {"id": "OpenFinAL/GPT2_FINGPT_QA", "local": True},
-    "InvestLM": {"id": "yixuantt/InvestLM-mistral-AWQ", "local": False},
-    "FinLLaMA": {"id": "us4/fin-llama3.1-8b", "local": False},
-    "FinanceConnect": {"id": "ceadar-ie/FinanceConnect-13B", "local": True},
-    "Sujet-Finance": {"id": "sujet-ai/Sujet-Finance-8B-v0.1", "local": True}
+    "FinGPT (GPT2)": {"id": "OpenFinAL/GPT2_FINGPT_QA", "local": True},
+    "InvestLM (AWQ)": {"id": "yixuantt/InvestLM-mistral-AWQ", "local": False},
+    "FinLLaMA (LLaMA3.1-8B)": {"id": "us4/fin-llama3.1-8b", "local": False},
+    "FinanceConnect (13B)": {"id": "ceadar-ie/FinanceConnect-13B", "local": True},
+    "Sujet-Finance (8B)": {"id": "sujet-ai/Sujet-Finance-8B-v0.1", "local": True}
 }
 
-# Load local model
+# === Load local models with caching ===
 @st.cache_resource
 def load_local_model(model_id):
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
@@ -31,27 +31,27 @@ def load_local_model(model_id):
     )
     return model, tokenizer
 
-# Build discursive prompt
+# === Build system prompt for discursive answers ===
 def build_prompt(user_question):
     return (
-        "You are a helpful and knowledgeable financial assistant named FinGPT. "
-        "You explain financial terms and concepts clearly, with examples when useful.\n\n"
+        "You are FinGPT, a helpful and knowledgeable financial assistant. "
+        "You explain finance, controlling, and tax topics clearly, with examples when useful.\n\n"
         f"User: {user_question.strip()}\n"
         "FinGPT:"
     )
 
-# Clean up repeated parts
+# === Clean repeated/extra outputs ===
 def clean_output(output_text):
     parts = output_text.split("FinGPT:")
     return parts[-1].strip() if len(parts) > 1 else output_text.strip()
 
-# Local inference
+# === Generate with local model ===
 def query_local_model(model_id, prompt):
     model, tokenizer = load_local_model(model_id)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=200,
+        max_new_tokens=300,
         temperature=0.7,
         top_k=50,
         top_p=0.95,
@@ -63,10 +63,10 @@ def query_local_model(model_id, prompt):
     raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return clean_output(raw_output)
 
-# Remote inference
+# === Generate with remote HF API ===
 def query_remote_model(model_id, prompt):
     headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 200}}
+    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 300}}
     response = requests.post(
         f"https://api-inference.huggingface.co/models/{model_id}",
         headers=headers,
@@ -76,31 +76,30 @@ def query_remote_model(model_id, prompt):
         result = response.json()
         return result[0]["generated_text"] if isinstance(result, list) else result.get("generated_text", "No output")
     else:
-        raise RuntimeError(f"API Error: {response.status_code} — {response.text}")
+        raise RuntimeError(f"API Error {response.status_code}: {response.text}")
 
-# Unified query handler
+# === Unified model query handler ===
 def query_model(model_entry, user_question):
     prompt = build_prompt(user_question)
     if model_entry["local"]:
         return query_local_model(model_entry["id"], prompt)
     else:
-        return query_remote_model(model_entry["id"], prompt)
+        return clean_output(query_remote_model(model_entry["id"], prompt))
 
-# Streamlit UI
-st.set_page_config(page_title="Financial LLM Interface", layout="centered")
+# === Streamlit UI Layout ===
+st.set_page_config(page_title="Finance LLM Comparison", layout="centered")
 st.title("💼 Financial LLM Evaluation Interface")
 
 model_choice = st.selectbox("Select a Financial Model", list(model_map.keys()))
-user_question = st.text_area("Enter your financial question:", "What is CAP in finance?")
+user_question = st.text_area("Enter your financial question:", "What is EBIT vs EBITDA?", height=150)
 
 if st.button("Get Response"):
-    with st.spinner("Generating discursive response..."):
+    with st.spinner("Thinking like a CFO..."):
         try:
             model_entry = model_map[model_choice]
             answer = query_model(model_entry, user_question)
-            st.markdown("### 🧠 Response:")
-            st.text_area("💬 Response from FinGPT:", value=answer, height=200, disabled=True)
-
+            st.text_area("💬 Response:", value=answer, height=300, disabled=True)
         except Exception as e:
             st.error(f"❌ Error: {e}")
 
+