Update inference.py

inference.py

@@ -1,26 +1,29 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-import torch.nn.functional as F
-from peft import PeftModel
-
-# Load model and tokenizer
-model_name = "munzirmuneer/phishing_url_gemma_pytorch"  # Replace with your specific model
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-model = PeftModel.from_pretrained(model, model_name)
-
-def predict(input_text):
-    # Tokenize input
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
-    
-    # Run inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Get logits and probabilities
-    logits = outputs.logits
-    probs = F.softmax(logits, dim=-1)
-    
-    # Get the predicted class (highest probability)
-    pred_class = torch.argmax(probs, dim=-1)
-    return pred_class.item(), probs[0].tolist()
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+from huggingface_hub import HfApi
+import torch.nn.functional as F
+from peft import PeftModel
+
+HfApi().set_access_token("HUGGINGFACE_HUB_TOKEN")
+
+# Load model and tokenizer
+model_name = "munzirmuneer/phishing_url_gemma_pytorch"  # Replace with your specific model
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
+model = AutoModelForSequenceClassification.from_pretrained(model_name, use_auth_token=True)
+model = PeftModel.from_pretrained(model, model_name, use_auth_token=True)
+
+def predict(input_text):
+    # Tokenize input
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+    
+    # Run inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Get logits and probabilities
+    logits = outputs.logits
+    probs = F.softmax(logits, dim=-1)
+    
+    # Get the predicted class (highest probability)
+    pred_class = torch.argmax(probs, dim=-1)
+    return pred_class.item(), probs[0].tolist()