Update app.py

app.py

@@ -1,9 +1,31 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+import os
 
-# Initialize the FastAPI app
 app = FastAPI()
 
-# Define the welcome endpoint
-@app.get('/')
-async def welcome():
-    return "Welcome to our Text Classification API"
+# 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", use_auth_token=token)
+tokenizer = AutoTokenizer.from_pretrained(
+    "kmcs-casulit/hr_cate", use_auth_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']}