Update app.py

app.py

@@ -5,10 +5,10 @@ from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
 import uvicorn
-import json
-from huggingface_hub import hf_hub_download, login
 
-# Authenticate with Hugging Face Hub
+from huggingface_hub import login
+
+# Authenticate with Hugging Face Hub using the HF_TOKEN environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
@@ -21,55 +21,23 @@ class Query(BaseModel):
 
 app = FastAPI(title="Financial Chatbot API")
 
-# Load the base model
-base_model_name = "meta-llama/Llama-3.2-3B"  # Update if using a different base model
-model = AutoModelForCausalLM.from_pretrained(
+# Load the base model from Meta-Llama
+base_model_name = "meta-llama/Llama-3.2-3B"
+base_model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
     device_map="auto",
     trust_remote_code=True
 )
 
-# Load adapter from your checkpoint with a fix for the 'eva_config' issue
+# Load the finetuned adapter using PEFT
 peft_model_id = "Phoenix21/llama-3-2-3b-finetuned-finance_checkpoint2"
+model = PeftModel.from_pretrained(base_model, peft_model_id)
 
-# Manually download and load the adapter config to filter out problematic fields
-try:
-    # Download the adapter_config.json file
-    config_file = hf_hub_download(
-        repo_id=peft_model_id,
-        filename="adapter_config.json",
-        token=HF_TOKEN
-    )
-    
-    # Load and clean the config
-    with open(config_file, 'r') as f:
-        config_dict = json.load(f)
-    
-    # Remove problematic fields if they exist
-    if "eva_config" in config_dict:
-        del config_dict["eva_config"]
-    
-    # Load the adapter directly with the cleaned config
-    model = PeftModel.from_pretrained(
-        model,
-        peft_model_id,
-        config=config_dict
-    )
-except Exception as e:
-    print(f"Error loading adapter: {e}")
-    # Fallback to direct loading if the above fails
-    model = PeftModel.from_pretrained(
-        model,
-        peft_model_id,
-        # Use this config parameter to ignore unknown parameters
-        config=None
-    )
-
-# Load tokenizer from the base model
+# Load the tokenizer from the base model
 tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-# Create a text-generation pipeline
+# Create a text-generation pipeline using the loaded model and tokenizer
 chat_pipe = pipeline(
     "text-generation",
     model=model,
@@ -83,10 +51,8 @@ chat_pipe = pipeline(
 def generate(query: Query):
     prompt = f"Question: {query.text}\nAnswer: "
     response = chat_pipe(prompt)[0]["generated_text"]
-    # Extract only the answer part from the response
-    answer = response.split("Answer: ")[-1].strip()
-    return {"response": answer}
+    return {"response": response}
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    uvicorn.run(app, host="0.0.0.0", port=port)