from fastapi import FastAPI, Query
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

app = FastAPI()

# Load the model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained(
    "kmcs-casulit/hr_cate")
tokenizer = AutoTokenizer.from_pretrained("kmcs-casulit/department-classification")
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)

# Define a request model
class TextRequest(BaseModel):
    text: str

@app.get("/")
def greet_json():
    return {"message": "Hello, World!"}


@app.post("/classify/")
def classify(request: TextRequest):
    output = pipe(request.text)
    return {"output": output[0]['label']}