Update app.py

app.py

@@ -3,10 +3,10 @@ import torch
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from peft import PeftModel, PeftConfig
+from peft import PeftModel
 import uvicorn
-
-from huggingface_hub import login
+import json
+from huggingface_hub import hf_hub_download, login
 
 # Authenticate with Hugging Face Hub
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -29,17 +29,41 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 
-# Load adapter from your checkpoint with a workaround for 'eva_config'
+# Load adapter from your checkpoint with a fix for the 'eva_config' issue
 peft_model_id = "Phoenix21/llama-3-2-3b-finetuned-finance_checkpoint2"
-# Load the PEFT configuration first
-peft_config = PeftConfig.from_pretrained(peft_model_id)
-# Remove 'eva_config' if it exists in the configuration
-peft_config_dict = peft_config.to_dict()
-if "eva_config" in peft_config_dict:
-    peft_config_dict.pop("eva_config")
-    peft_config = PeftConfig.from_dict(peft_config_dict)
-# Load the adapter using the filtered configuration
-model = PeftModel.from_pretrained(model, peft_model_id, config=peft_config)
+
+# 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
 tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
@@ -59,8 +83,10 @@ chat_pipe = pipeline(
 def generate(query: Query):
     prompt = f"Question: {query.text}\nAnswer: "
     response = chat_pipe(prompt)[0]["generated_text"]
-    return {"response": response}
+    # Extract only the answer part from the response
+    answer = response.split("Answer: ")[-1].strip()
+    return {"response": answer}
 
 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)