Update app.py

app.py

@@ -4,14 +4,14 @@ 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
+from preprocess_data import remove_html, remove_links
 import gradio as gr
 
+# Define input model structure
 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)
@@ -24,7 +24,7 @@ for param in bert.parameters():
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load custom model
+# Load your custom BERT-based model
 model = BERT_Arch(bert)
 model.load_state_dict(torch.load(model_path, map_location=device))
 model.to(device)
@@ -42,36 +42,37 @@ def home():
 
 @app.post("/predict/")
 def predict(request: TextRequest):
-    try:
-        text = request.text.strip()
+    return {"prediction": classify_text(request.text)}
 
-        # Preprocess text
+# Function to classify text
+def classify_text(text: str) -> str:
+    try:
+        text = text.strip()
         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
+            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"}
+        return "Phishing" if prediction == 1 else "Not Phishing"
 
     except Exception as e:
-        return {"error": str(e)}
-def greet(name):
-    return "Hello " + name + "!"
+        return f"Error: {str(e)}"
+
+# Gradio UI
 gr.Interface(
-    fn=greet,
-    inputs="text",
-    outputs="text",
-    allow_flagging="never"
-).launch()
+    fn=classify_text,
+    inputs=gr.Textbox(label="Enter website content or email text"),
+    outputs=gr.Label(label="Prediction"),
+    title="Phishing Text Detector",
+    description="Website text to check if it's phishing."
+).launch()