Update app.py

app.py

@@ -1,25 +1,50 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import numpy as np
+import torch
 
-# Initialize the FastAPI app
 app = FastAPI()
 
-# Initialize the model pipeline
-pipe = pipeline("text-classification", model="kmack/malicious-url-detection")
+# Check if CUDA is available
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+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
+
 # Define the API endpoint for URL prediction
 @app.post("/predict")
 async def predict(url_request: URLRequest):
     url_to_check = url_request.url
-    result = pipe(url_to_check)
+    result = get_prediction(url_to_check)
     return {"prediction": result}
 
 # Health check endpoint
 @app.get("/")
 async def read_root():
     return {"message": "API is up and running"}
+