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()