File size: 2,191 Bytes
96a8ab7 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import gdown
import torch
from fastapi import FastAPI
from transformers import AutoModel, BertTokenizerFast
from pydantic import BaseModel
from model import BERT_Arch
from preprocess_data import remove_html,remove_links
import gradio as gr
class TextRequest(BaseModel):
text: str
# Download model from Google Drive
#link:https://drive.google.com/drive/folders/102UPd446eHCCENR58EC3UxnJfcYkBa8U?usp=sharing
model_url = "https://drive.google.com/uc?id=16ZWVa0d2V0T3s11Oq86rLOTA6bOR0DnR"
model_path = "model.pth"
gdown.download(model_url, model_path, quiet=False)
# Load pre-trained BERT model
bert = AutoModel.from_pretrained("bert-base-uncased")
for param in bert.parameters():
param.requires_grad = False
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load custom model
model = BERT_Arch(bert)
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device)
model.eval()
# Load tokenizer
tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
# Initialize FastAPI
app = FastAPI()
@app.get("/")
def home():
return {"message": "Phishing Detection API is running!"}
@app.post("/predict/")
def predict(request: TextRequest):
try:
text = request.text.strip()
# Preprocess text
text = remove_html(text)
text = remove_links(text)
# Tokenize input text
tokens = tokenizer(
text, return_tensors="pt", truncation=True, padding="max_length", max_length=512
)
input_ids = tokens["input_ids"].to(device)
attention_mask = tokens["attention_mask"].to(device)
# Perform inference
with torch.no_grad():
output = model(input_ids, attention_mask)
prediction = torch.argmax(output.cpu(), dim=1).item()
return {"prediction": "Phishing" if prediction == 1 else "Not Phishing"}
except Exception as e:
return {"error": str(e)}
def greet(name):
return "Hello " + name + "!"
gr.Interface(
fn=greet,
inputs="text",
outputs="text",
allow_flagging="never"
).launch() |