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

# Define the input schema
class InputData(BaseModel):
    input: str

# Initialize the pipeline
pipe = pipeline("text-classification", model="phishbot/ScamLLM")

app = FastAPI()

# Define API endpoints
@app.post("/infer")
async def infer(data: InputData):
    predictions = pipe(data.input)
    return predictions

@app.get("/health")
async def health():
    return {"message": "ok"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)