Urna eletrônica com Python - eleições para a Prefeitura de São Paulo

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import pandas as pd

# Carregando o modelo e o tokenizer do DialoGPT
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# Dados dos candidatos para votação fictícia
dados_candidatos = {
    "numeros_candidatos": ["16", "27", "45", "50", "29", "30", "28", "15", "80", "40"],
    "nomes": ["Altino Prazeres (PSTU)", "Bebeto Haddad (DC)", "Datena (PSDB)", "Guilherme Boulos (PSOL)", 
             "João Pimenta (PCO)", "Marina Helena (NOVO)", "Pablo Marçal (PRTB)", "Ricardo Nunes (MDB)", "Ricardo Senese (UP)", "Tabata Amaral (PSB)"]
}
# Criando um DataFrame com os dados dos candidatos
df_dados_candidatos = pd.DataFrame(dados_candidatos)
df_dados_candidatos

# Função para verificar se o número do candidato é válido e retornar o nome
def verificar_candidato(numero_pedido):
    try:
        # Buscando o nome do candidato com base no número fornecido
        nome_candidato = df_dados_candidatos[df_dados_candidatos["numeros_candidatos"] == numero_pedido]["nomes"].iloc[0]
        return f"You are voting on {nome_candidato} that has the number: {numero_pedido}"
    except:
        # Caso não encontre o candidato, retorna uma mensagem de voto nulo
        return "Null vote!"

# Lista de palavras-chave que indicam a intenção de votar
palavras_chave_voto = ["choice", "vote", "choose"]

# Histórico de conversa inicializado como None
ids_historico_chat = None

while True:
    input_usuario = input('You: ')

    # Verificando confirmação para encerrar o loop
    if input_usuario.lower() in ["confirm", "yes", "ok", "sure", "that's ok"]:
        print('Bot: Congratulations on exercising your citizenship!')
        break

    # Se o usuário menciona palavras relacionadas ao voto, inicia o processo de votação
    if any(keyword in input_usuario.lower() for keyword in palavras_chave_voto):
        # Pergunta ao usuário qual número ele deseja votar
        numero_candidato = input('''Bot: Choose one of the numbers below to take over as mayor of São Paulo: 
        16
        27
        45
        50
        29
        30
        28
        15
        80
        40
        
        ''')
        # Verifica se o número é válido e retorna o nome do candidato ou voto nulo
        resposta = verificar_candidato(numero_candidato)
        print('Bot: Do you confirm your vote?')
    
    # Se o input não for relacionado ao voto, responde usando o modelo DialoGPT
    else:
        # Tokenizando a entrada do usuário e adicionando o token especial EOS (fim de sequência)
        novo_usuario_input_ids = tokenizer.encode(input_usuario + tokenizer.eos_token, return_tensors='pt')
        
        # Se já houver histórico de conversa, concatenar a nova entrada ao histórico
        if ids_historico_chat is not None:
            bot_input_ids = torch.cat([ids_historico_chat, novo_usuario_input_ids], dim=-1)
        else:
            # Caso contrário, usar apenas a nova entrada
            bot_input_ids = novo_usuario_input_ids
            
        # Gerar a resposta do bot usando o modelo
        ids_historico_chat = model.generate(
            bot_input_ids,
            max_length=1000,
            pad_token_id=tokenizer.eos_token_id
        )
        # Decodifica a resposta gerada para string
        resposta = tokenizer.decode(ids_historico_chat[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    
    # Imprime a resposta do bot
    print(f'Bot: {resposta}')