Upload inference.py

inference.py

@@ -0,0 +1,26 @@
+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()