import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer
import json

# Inicialize o cliente e o tokenizador
model_name = "rss9051/autotrein-BERT-iiLEX-dgs-0004"
client = InferenceClient(model=model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Função para dividir o texto em chunks menores com o tokenizador
def split_text_into_chunks(text, max_tokens=512):
    tokens = tokenizer(text, return_tensors="pt", truncation=False)["input_ids"][0]
    chunks = []
    for i in range(0, len(tokens), max_tokens):
        chunk = tokens[i:i + max_tokens]
        chunks.append(chunk)
    return [tokenizer.decode(chunk, skip_special_tokens=True) for chunk in chunks]

# Função para classificar texto longo
def classify_text(text):
    chunks = split_text_into_chunks(text, max_tokens=512)  # Divida o texto em chunks de 512 tokens
    all_responses = []  # Lista para armazenar respostas de cada chunk

    for chunk in chunks:
        response_bytes = client.post(json={"inputs": chunk})  # Enviar o chunk
        response_str = response_bytes.decode('utf-8')  # Decodificar de bytes para string
        response = json.loads(response_str)  # Converter string JSON para objeto Python

        if isinstance(response, list) and len(response) > 0:
            sorted_response = sorted(response[0], key=lambda x: x['score'], reverse=True)
            all_responses.append(sorted_response[0])  # Adicionar a melhor classificação do chunk

    # Combinar resultados de todos os chunks
    if all_responses:
        # Contar as classes mais frequentes
        class_scores = {}
        for res in all_responses:
            label = res['label']
            score = res['score']
            if label in class_scores:
                class_scores[label] += score
            else:
                class_scores[label] = score
        
        # Obter a classe com maior score combinado
        predicted_class = max(class_scores, key=class_scores.get)
    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()