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

app = FastAPI()

# Set the Hugging Face token from the environment variable
token = os.getenv("token")

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


class TextInput(BaseModel):
    text: str


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


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