app.py

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

app = FastAPI()
token = os.getenv("token")

# Load the model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained(
    "kmcs-casulit/hr_cate", use_auth_token=token)
tokenizer = AutoTokenizer.from_pretrained("kmcs-casulit/department-classification", use_auth_token=token)
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']}