Update inference.py

inference.py

@@ -39,9 +39,12 @@ class MilUnixCoder(nn.Module):
 
 def load_model():
     """Load the model and tokenizer"""
-    tokenizer = AutoTokenizer.from_pretrained("YoungDSMLKZ/UnixCoder-MIL")
+    from huggingface_hub import hf_hub_download
+    repo = "YoungDSMLKZ/UnixCoder-MIL"
+    tokenizer = AutoTokenizer.from_pretrained(repo)
     model = MilUnixCoder("microsoft/unixcoder-base")
-    model.load_state_dict(load_file("YoungDSMLKZ/UnixCoder-MIL/model.safetensors"))
+    weights_path = hf_hub_download(repo_id=repo, filename="model.safetensors")
+    model.load_state_dict(load_file(weights_path))
     model.to(DEVICE).eval()
     return model, tokenizer