app.py

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

# Initialize the FastAPI app
app = FastAPI()

# Initialize the model pipeline
pipe = pipeline("text-classification", model="kmack/malicious-url-detection")

# Define the request model
class URLRequest(BaseModel):
    url: str

# 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)
    return {"prediction": result}

# Health check endpoint
@app.get("/")
async def read_root():
    return {"message": "API is up and running"}