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

# Load the depression detection model
tokenizer = AutoTokenizer.from_pretrained("paulagarciaserrano/roberta-depression-detection")
model = AutoModelForSequenceClassification.from_pretrained("paulagarciaserrano/roberta-depression-detection")
depression_pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)

# Pydantic object for validation
class DepressionInput(BaseModel):
    text: str

# FastAPI app
app = FastAPI()

# Endpoint for depression detection
@app.post("/depression_detection")
async def detect_depression(item: DepressionInput):
    result = depression_pipe(item.text)
    return result