Update app.py

app.py

@@ -8,7 +8,7 @@
 
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
 
 app = FastAPI()
@@ -19,29 +19,17 @@ if torch.cuda.is_available():
 else:
     device = torch.device("cpu")
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("kmack/malicious-url-detection")
-model = AutoModelForSequenceClassification.from_pretrained("kmack/malicious-url-detection")
-model = model.to(device)
-
 # Define the request model
 class URLRequest(BaseModel):
     url: str
 
-# Prediction function
-def get_prediction(input_text: str) -> dict:
-    label2id = model.config.label2id
-    inputs = tokenizer(input_text, return_tensors='pt', truncation=True)
-    inputs = inputs.to(device)
-    outputs = model(**inputs)
-    logits = outputs.logits
-    sigmoid = torch.nn.Sigmoid()
-    probs = sigmoid(logits.squeeze().cpu())
-    probs = probs.detach().numpy()
-    for i, k in enumerate(label2id.keys()):
-        label2id[k] = probs[i]
-    label2id = {k: float(v) for k, v in sorted(label2id.items(), key=lambda item: item[1].item(), reverse=True)}
-    return label2id
+# Load the tokenizer and model using pipeline
+pipe = pipeline("text-classification", model="kmack/malicious-url-detection", device=device.index if torch.cuda.is_available() else -1)
+
+# Define the prediction function
+def get_prediction(url_to_check: str):
+    result = pipe(url_to_check)
+    return result
 
 # Define the API endpoint for URL prediction
 @app.post("/predict")
@@ -54,3 +42,4 @@ async def predict(url_request: URLRequest):
 @app.get("/")
 async def read_root():
     return {"message": "API is up and running"}
+