from transformers import pipeline

# Pipeline do modelo BART
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Dicionário com descrições das categorias como chave e descrições como valor.
# Traz a vantagem de conter mais contexto semântico para o modelo.
# Pode melhorar a discriminação entre categorias próximas (ex: política vs. economia).

categorias_descritas = {
    "esporte": "notícias relacionadas a competições esportivas, atletas ou eventos esportivos",
    "política": "notícias sobre decisões governamentais, eleições, partidos ou políticas públicas",
    "tecnologia": "notícias sobre inovações, ciência da computação, internet e dispositivos eletrônicos",
    "entretenimento": "notícias sobre filmes, música, celebridades e cultura pop", # nova categoria
    "meio ambiente": "notícias sobre mudanças climáticas, preservação ambiental e ecologia",
    "saúde": "notícias sobre medicina, doenças, vacinas e cuidados com a saúde"
}

# Lista de descrições para o modelo
descricoes = list(categorias_descritas.values())

# Função de predição
def classificar_artigo(texto):
    resultado = classifier(texto, candidate_labels=descricoes)

    # Converte a descrição de volta para o nome da categoria original
    saida = {}
    for desc, score in zip(resultado["labels"], resultado["scores"]):
        for categoria, descricao in categorias_descritas.items():
            if desc == descricao:
                saida[categoria] = f"{score:.4f}"
                break
    return saida

# Interface Gradio
import gradio as gr
interface = gr.Interface(
    fn=classificar_artigo,
    inputs=gr.Textbox(lines=10, label="Digite o conteúdo do artigo"),
    outputs=gr.Label(label="Classificação Zero-Shot"),
    title="Classificação Zero-Shot de notícias",
    description="Este modelo classifica notícias com base em descrições detalhadas das categorias."
)

interface.launch()