import gradio as gr
from huggingface_hub import InferenceClient
import json

# Inicialize o cliente com o modelo do Hugging Face
#client = InferenceClient(model="ulisesbravo/autotrain-nsuej-5ctie")
#client = InferenceClient(model="ulisesbravo/autotrain-nzog3-ca819");
#client = InferenceClient(model="ulisesbravo/autotrain-po0st-um4bf");
#client = InferenceClient(model="ulisesbravo/autotrain-k9ag8-b7nm0");
#client = InferenceClient(model="ulisesbravo/autotrain-pt-l1sfd-nom8j");
#client = InferenceClient(model="ulisesbravo/autotrain-p3-portuguese-sdbvu-hn7p7");
client = InferenceClient(model="ulisesbravo/autotrain-BERTimbau-DGS-0003");

def classify_text(text):
    # Realize a inferência chamando o método post
    response_bytes = client.post(json={"inputs": text})  # Enviar o texto

    # Decodificar a resposta de bytes para string e depois para JSON
    response_str = response_bytes.decode('utf-8')  # Decodificar de bytes para string
    response = json.loads(response_str)  # Converter string JSON para um objeto Python

    # Inspecionar a resposta para depuração
    print(response)

    # Verificar se a resposta é uma lista válida
    if isinstance(response, list) and len(response) > 0:
        # Ordenar as classificações pelo score e pegar a de maior valor
        sorted_response = sorted(response[0], key=lambda x: x['score'], reverse=True)
        predicted_class = sorted_response[0]['label']  # Pegar a classe com maior score
    else:
        predicted_class = "Classificação não encontrada"

    return predicted_class

# Interface Gradio
demo = gr.Interface(
    fn=classify_text,  # Função a ser chamada para classificar o texto
    inputs=gr.Textbox(label="Texto para Classificação"),  # Entrada de texto
    outputs=gr.Label(label="Classe Predita"),  # Saída da classificação
    title="Classificador de Texto",  # Título da interface
    description="Insira um texto para obter a classificação usando o modelo treinado."  # Descrição da interface
)

if __name__ == "__main__":
    demo.launch()