shahzaib201's picture
Update main.py
8bcd317 verified
raw
history blame contribute delete
766 Bytes
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