from fastapi import FastAPI
from transformers import pipeline

app=FastAPI()

pipe = pipeline("zero-shot-classification", model="joeddav/xlm-roberta-large-xnli")

@app.get("/")
def home():
    return {"message": "It is working!"}

@app.get("/generate")
def generate(sequence_to_classify):
    classifier = pipeline("zero-shot-classification", model="joeddav/xlm-roberta-large-xnli")

    candidate_labels = ["SAD", "HAPPY", "AGGRESSIVE"]

    output = classifier(sequence_to_classify, candidate_labels)

    return {"output":output['labels'][0]}