Update app.py

app.py

@@ -1,35 +1,66 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
 import torch
+import requests
 import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import login, HfApi
 
+# Optional: Login if you want access to gated/private models
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+if HF_TOKEN:
+    login(HF_TOKEN)
 
-# Set model map
+# Define model map with access type
 model_map = {
-    "FinGPT": "AI4Finance/FinGPT",
-    "FinanceConnect": "ceadar-ie/FinanceConnect-13B",
-    "Sujet-Finance": "sujet-ai/Sujet-Finance-8B-v0.1"
+    "FinGPT": {"id": "AI4Finance/FinGPT", "local": True},
+    "InvestLM": {"id": "mrm8488/investLM-7B", "local": False},  # example ID, update if needed
+    "FinLLaMA": {"id": "HuggingFaceH4/fin-llama", "local": False},
+    "FinanceConnect": {"id": "ceadar-ie/FinanceConnect-13B", "local": True},
+    "Sujet-Finance": {"id": "sujet-ai/Sujet-Finance-8B-v0.1", "local": True}
 }
 
-# Cache model loading for performance
+# Cache local models
 @st.cache_resource
-def load_model_and_tokenizer(model_id):
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+def load_local_model(model_id):
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        use_auth_token=HF_TOKEN
     )
     return model, tokenizer
 
-# Query model
-def query_model(model_id, question):
-    model, tokenizer = load_model_and_tokenizer(model_id)
-    inputs = tokenizer(question, return_tensors="pt")
+# Local model querying
+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=150)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Streamlit app layout
+# Remote model querying (via Inference 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": 150}}
+    response = requests.post(
+        f"https://api-inference.huggingface.co/models/{model_id}",
+        headers=headers,
+        json=payload
+    )
+    if response.status_code == 200:
+        result = response.json()
+        return result[0]["generated_text"] if isinstance(result, list) else result.get("generated_text", "No output")
+    else:
+        raise RuntimeError(f"Failed to call remote model: {response.text}")
+
+# Unified query dispatcher
+def query_model(model_entry, prompt):
+    if model_entry["local"]:
+        return query_local_model(model_entry["id"], prompt)
+    else:
+        return query_remote_model(model_entry["id"], prompt)
+
+# --- Streamlit UI ---
 st.title("💼 Financial LLM Evaluation Interface")
 
 model_choice = st.selectbox("Select a Financial Model", list(model_map.keys()))
@@ -38,7 +69,8 @@ user_question = st.text_area("Enter your financial question:", "What is EBITDA?"
 if st.button("Get Response"):
     with st.spinner("Generating response..."):
         try:
-            answer = query_model(model_map[model_choice], user_question)
+            model_entry = model_map[model_choice]
+            answer = query_model(model_entry, user_question)
             st.subheader(f"Response from {model_choice}:")
             st.write(answer)
         except Exception as e: