Update app.py

app.py

@@ -5,6 +5,8 @@ import autopep8
 import glob
 import re
 import os
+from huggingface_hub import hf_hub_download
+
 
 # ==========================
 #  Utility functions
@@ -188,17 +190,30 @@ def find_example_files():
 #  Load model from HF Hub
 # ==========================
 
-BASE_MODEL_ID = "Salesforce/codet5p-770m"             
-FINETUNED_MODEL_ID = "OSS-Forge/codet5p-770m-pyresbugs"
+BASE_MODEL_ID = "Salesforce/codet5p-770m"
+FINETUNED_REPO_ID = "OSS-forge/codet5p-770m-pyresbugs"   
+FINETUNED_FILENAME = "pytorch_model.bin"
 
 print(f"Loading tokenizer from base model: {BASE_MODEL_ID}")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
 
-print(f"Loading fine-tuned model weights from: {FINETUNED_MODEL_ID}")
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    FINETUNED_MODEL_ID,
-    ignore_mismatched_sizes=True,   
-)
+print(f"Loading base model: {BASE_MODEL_ID}")
+model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL_ID)
+
+print(f"Downloading fine-tuned weights from repo: {FINETUNED_REPO_ID}")
+ckpt_path = hf_hub_download(FINETUNED_REPO_ID, FINETUNED_FILENAME)
+
+print(f"Loading state_dict from: {ckpt_path}")
+state_dict = torch.load(ckpt_path, map_location=torch.device("cpu"))
+
+if "model_state_dict" in state_dict:
+    state_dict = state_dict["model_state_dict"]
+
+missing, unexpected = model.load_state_dict(state_dict, strict=False)
+print("Loaded fine-tuned weights.")
+print("Missing keys:", len(missing), "Unexpected keys:", len(unexpected))
+
+model.eval()