File size: 1,724 Bytes
374ec5a fdad04a 374ec5a b0227dc fdad04a b0227dc fdad04a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
from peft import PeftModel
class EndpointHandler:
def __init__(self, model_dir):
"""
Initialize the model and tokenizer using the provided model directory.
"""
model_name = "munzirmuneer/phishing_url_gemma_pytorch" # Replace with your specific model
model_name2 = "google/gemma-2b"
# Load tokenizer and model
self.tokenizer = AutoTokenizer.from_pretrained(model_name2)
base_model = AutoModelForSequenceClassification.from_pretrained(model_name)
self.model = PeftModel.from_pretrained(base_model, model_name)
def __call__(self, input_data):
"""
Perform inference on the input text and return predictions.
"""
# Extract the URL from the input_data dictionary
if 'inputs' in input_data:
input_text = input_data['inputs'] # Expecting a single URL as a string
else:
raise ValueError("Input data must contain the 'inputs' key with a URL.")
# Tokenize input
inputs = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
# Run inference
with torch.no_grad():
outputs = self.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 {
"predicted_class": pred_class.item(),
"probabilities": probs[0].tolist()
}
|