Update app.py

app.py

@@ -3,26 +3,31 @@ import torch
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from peft import PeftModel, PeftConfig
 import uvicorn
 
-# Define a Pydantic model for request validation
 class Query(BaseModel):
     text: str
 
-# Initialize FastAPI app
 app = FastAPI(title="Financial Chatbot API")
 
-# Load your fine-tuned model and tokenizer using the updated model name
-model_name = "Phoenix21/meta-llama-Llama-3.2-3B-2025-03-13-checkpoints"
+# Load base model
+base_model_name = "meta-llama/Meta-Llama-3-8B"  # Update this if different base model
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
+    base_model_name,
     device_map="auto",
     trust_remote_code=True
 )
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+# Load adapter from your checkpoint
+peft_model_id = "Phoenix21/llama-3-2-3b-finetuned-finance_checkpoint2"
+model = PeftModel.from_pretrained(model, peft_model_id)
+
+# Load tokenizer from base model
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-# Create a text-generation pipeline
+# Rest of your code remains the same...
 chat_pipe = pipeline(
     "text-generation",
     model=model,
@@ -32,14 +37,12 @@ chat_pipe = pipeline(
     top_p=0.95,
 )
 
-# Define an endpoint for generating responses
 @app.post("/generate")
 def generate(query: Query):
     prompt = f"Question: {query.text}\nAnswer: "
     response = chat_pipe(prompt)[0]["generated_text"]
     return {"response": response}
 
-# Run the app using uvicorn; default port is 7860 (as expected by Hugging Face Spaces)
 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)