File size: 569 Bytes
151ed35
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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)