File size: 1,982 Bytes
103e422 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 | from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
app = FastAPI()
# NEW MODEL: Multimodal Phishing Detector (URLs, SMS, Email)
MODEL_ID = "ealvaradob/bert-finetuned-phishing"
print("Loading model... This might take a minute as it's a 'large' BERT model.")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
class URLInput(BaseModel):
url: str
@app.get("/")
async def root():
return {"status": "URL Phishing Detector API is running"}
@app.post("/predict")
async def predict_url(data: URLInput):
# 1. Basic Pre-check
if not data.url or len(data.url) < 4:
return {"error": "Invalid URL provided"}
# 2. Tokenize and Predict
inputs = tokenizer(data.url, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
# Apply Softmax to get percentages
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
probs = predictions[0].tolist()
# 3. Dynamic Label Mapping
# The model usually uses LABEL_0 (Legitimate) and LABEL_1 (Phishing)
confidences = {model.config.id2label[i]: prob for i, prob in enumerate(probs)}
# Identify the highest confidence label
max_label = max(confidences.items(), key=lambda x: x[1])
label_name = max_label[0]
# Check for "LABEL_1" or "phishing" keyword in the output
is_phishing = "1" in label_name or "phishing" in label_name.lower()
return {
"url": data.url,
"prediction": "phishing" if is_phishing else "legitimate",
"confidence": round(max_label[1], 4),
"raw_scores": confidences,
"is_malicious": is_phishing
}
if __name__ == "__main__":
import uvicorn
# 7860 is the standard port for Hugging Face Spaces
uvicorn.run(app, host="0.0.0.0", port=7860) |