Update app.py

app.py

@@ -4,6 +4,26 @@ import requests
 import os
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+import torch
+
+@st.cache_resource
+def load_fingpt_lora():
+    base_model_id = "meta-llama/Llama-2-7b-hf"
+    lora_adapter_id = "FinGPT/fingpt-mt_llama2-7b_lora"
+    tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_auth_token=HF_TOKEN)
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto",
+        use_auth_token=HF_TOKEN
+    )
+
+    model = PeftModel.from_pretrained(base_model, lora_adapter_id, use_auth_token=HF_TOKEN)
+    return model, tokenizer
+
 
 # Load token from Hugging Face Space secrets
 HF_TOKEN = os.getenv("Allie", None)
@@ -12,7 +32,7 @@ if HF_TOKEN:
 
 # === Available Models for Selection ===
 model_map = {
-    "FinGPT (GPT2)": {"id": "OpenFinAL/GPT2_FINGPT_QA", "local": True},
+    "FinGPT LoRA" = {"id": "FinGPT/fingpt-mt_llama2-7b_lora", "local": True, "custom_loader": load_fingpt_lora},
     "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},