Spaces:
Sleeping
Sleeping
File size: 1,154 Bytes
46122a7 609af13 46122a7 5979505 46122a7 609af13 46122a7 5979505 46122a7 609af13 |
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 |
# from fastapi import FastAPI
# app = FastAPI()
# @app.get("/")
# def greet_json():
# return {"Hello": "World!"}
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch
app = FastAPI()
# Check if CUDA is available
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
# Define the request model
class URLRequest(BaseModel):
url: str
# 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")
async def predict(url_request: URLRequest):
url_to_check = url_request.url
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"}
|