import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, pipeline import torch import numpy as np # === CARREGAR OS MODELOS GERADORES === generator_1_name = "pierreguillou/gpt2-small-portuguese" generator_2_name = "pierreguillou/gpt2-small-portuguese" # Usando o mesmo por simplicidade/teste tokenizer_1 = AutoTokenizer.from_pretrained(generator_1_name) model_1 = AutoModelForCausalLM.from_pretrained(generator_1_name) tokenizer_2 = AutoTokenizer.from_pretrained(generator_2_name) model_2 = AutoModelForCausalLM.from_pretrained(generator_2_name) # === CARREGAR MODELO ÁRBITRO (BERT) === judge_model_name = "neuralmind/bert-base-portuguese-cased" judge_tokenizer = AutoTokenizer.from_pretrained(judge_model_name) judge_model = AutoModelForSequenceClassification.from_pretrained(judge_model_name, num_labels=2) # Classificador de similaridade (baseado em relevância para o prompt) def score_response(prompt, response): inputs = judge_tokenizer(prompt, response, return_tensors="pt", truncation=True, padding=True) with torch.no_grad(): outputs = judge_model(**inputs) score = torch.softmax(outputs.logits, dim=1)[0][1].item() # Probabilidade da classe "boa" return score # Gerar resposta com modelo def generate_response(model, tokenizer, prompt): input_ids = tokenizer.encode(prompt, return_tensors="pt") output_ids = model.generate(input_ids, max_new_tokens=60, num_return_sequences=1, do_sample=True) return tokenizer.decode(output_ids[0], skip_special_tokens=True) # Função principal def chatbot(prompt): response_1 = generate_response(model_1, tokenizer_1, prompt) response_2 = generate_response(model_2, tokenizer_2, prompt) score_1 = score_response(prompt, response_1) score_2 = score_response(prompt, response_2) if score_1 > score_2: final = response_1 chosen = "Resposta 1" else: final = response_2 chosen = "Resposta 2" return ( prompt, response_1, response_2, chosen, final ) # === INTERFACE GRADIO === iface = gr.Interface( fn=chatbot, inputs=gr.Textbox(label="Digite sua pergunta"), outputs=[ gr.Textbox(label="Prompt"), gr.Textbox(label="Resposta 1"), gr.Textbox(label="Resposta 2"), gr.Textbox(label="Resposta escolhida pelo árbitro"), gr.Textbox(label="Resposta final exibida") ], title="Chatbot em Cascata (Português)", description="Dois modelos geram respostas e um árbitro (BERT) escolhe a melhor." ) if __name__ == "__main__": iface.launch()