scorer.py

#!/usr/bin/env python3
import sys
import os

# Passo 1: Parsear os argumentos de linha de comando antecipadamente para definir variáveis de ambiente
# Antes de importar transformers, parseamos os argumentos que podem afetar a configuração do modelo
def early_parse_args():
    import argparse

    parser = argparse.ArgumentParser(description="Scorer de Qualidade para Datasets do Hugging Face (Early Parsing)")
    parser.add_argument(
        "--4bit",
        action="store_true",
        dest="fourbit",
        help="Carrega o modelo em precisão de 4 bits utilizando BitsAndBytes."
    )
    # Parse apenas os argumentos necessários para BitsAndBytes
    args, _ = parser.parse_known_args()
    return args

early_args = early_parse_args()

# Definir a variável de ambiente antes de importar transformers, se --4bit for usado
if early_args.fourbit:
    os.environ["LLM_INT8_ENABLE_FP32_CPU_OFFLOAD"] = "true"

# Passo 2: Importar as bibliotecas restantes após definir variáveis de ambiente
import argparse
import json
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import numpy as np
from scipy.special import softmax
from datasets import load_dataset
import pyarrow.parquet as pq
from tqdm import tqdm
import torch

def infer_quality(model, tokenizer, input_text, resp_text, device):
    """
    Calcula a pontuação de qualidade para uma resposta baseada na entrada e na resposta fornecida.
    
    Args:
        model: O modelo de linguagem causal carregado.
        tokenizer: O tokenizer correspondente ao modelo.
        input_text (str): O texto da instrução/pergunta.
        resp_text (str): O texto da resposta a ser avaliada.
        device: O dispositivo (CPU ou GPU) para realizar os cálculos.
    
    Returns:
        float: A pontuação de qualidade calculada.
    """
    # Template traduzido para português
    quality_template = (
        "Você é um assistente útil. Por favor, identifique a pontuação de qualidade da Resposta correspondente à Pergunta.\n"
        "#Pergunta#:\n{instruction}\n"
        "#Resposta#:\n{output}\n"
        "##Qualidade: "
    )
    
    # Formatar o input do usuário
    user_input = quality_template.format(instruction=input_text, output=resp_text)
    
    # Tokenizar a entrada
    input_ids = tokenizer.encode(user_input, return_tensors="pt").to(device)
    max_length = 512  # Definir comprimento máximo
    
    # Gerar a resposta do modelo
    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_length=max_length,
            num_return_sequences=1,
            return_dict_in_generate=True,
            output_scores=True,
            do_sample=False  # Desativar amostragem para obter logits determinísticos
        )
    
    # Extrair os scores das saídas
    scores = outputs.scores  # Tuple of tensors, one for each generated token
    score_logits = []
    
    # Mapeamento de IDs de tokens para pontuações
    id2score = {
        29896: "1",  # Token ID para "1"
        29906: "2",  # Token ID para "2"
        29941: "3",  # Token ID para "3"
        29946: "4",  # Token ID para "4"
        29945: "5",  # Token ID para "5"
        29953: "6"   # Token ID para "6"
    }
    
    score_template = np.array([1, 2, 3, 4, 5, 6])
    
    # Verificar se há scores gerados
    if not scores:
        return 0.0  # Retorna 0 se não houver scores
    
    # Assumindo que queremos a última posição gerada
    last_score = scores[-1]  # Último conjunto de logits, shape: (batch_size * num_beams, vocab_size)
    
    # Garantir que last_score tenha o formato esperado
    if last_score.ndim != 2 or last_score.size(0) != 1:
        print("Formato inesperado para 'last_score'.")
        return 0.0
    
    last_score = last_score[0]  # Shape: (vocab_size,)
    
    # Extrair logits para os token_ids específicos
    for token_id in id2score:
        if token_id < last_score.size(0):
            logit = last_score[token_id].item()
            score_logits.append(logit)
        else:
            # Se o token_id estiver fora do alcance, atribuir um valor muito baixo
            score_logits.append(-1e10)
    
    score_logits = np.array(score_logits)
    
    # Aplicar softmax para obter probabilidades
    score_probs = softmax(score_logits)
    
    # Calcular a pontuação ponderada
    quality_score = np.sum(score_probs * score_template)
    
    return quality_score

def process_dataset(model, tokenizer, dataset_repo, input_field, output_field, output_dir, device):
    """
    Processa um dataset específico, calculando a pontuação de qualidade para cada exemplo e salvando o resultado.
    
    Args:
        model: O modelo de linguagem causal carregado.
        tokenizer: O tokenizer correspondente ao modelo.
        dataset_repo (str): O repositório do dataset no Hugging Face.
        input_field (str): O nome do campo de entrada no dataset.
        output_field (str): O nome do campo de resposta no dataset.
        output_dir (str): O diretório onde os datasets processados serão salvos.
        device: O dispositivo (CPU ou GPU) para realizar os cálculos.
    """
    print(f"\nProcessando dataset: {dataset_repo}")
    try:
        dataset = load_dataset(dataset_repo)
    except Exception as e:
        print(f"Erro ao carregar o dataset {dataset_repo}: {e}")
        return
    
    # Considerar apenas splits que contêm dados (por exemplo, 'train', 'validation', 'test')
    for split in dataset.keys():
        split_data = dataset[split]
        if len(split_data) == 0:
            continue
        
        print(f"  Processando split: {split} com {len(split_data)} exemplos")
        
        # Preparar listas para armazenar os scores
        quality_scores = []
        
        # Iterar sobre os exemplos com tqdm para visualizar o progresso
        for example in tqdm(split_data, desc=f"    Avaliando {split}"):
            input_text = example.get(input_field, "")
            output_text = example.get(output_field, "")
            if not isinstance(input_text, str) or not isinstance(output_text, str):
                quality_scores.append(0.0)
                continue
            score = infer_quality(model, tokenizer, input_text, output_text, device)
            quality_scores.append(score)
        
        # Adicionar a nova coluna ao dataset
        split_data = split_data.add_column("quality_score", quality_scores)
        
        # Definir o caminho de salvamento
        dataset_name = dataset_repo.split('/')[-1]
        split_output_dir = os.path.join(output_dir, dataset_name)
        os.makedirs(split_output_dir, exist_ok=True)
        output_path = os.path.join(split_output_dir, f"{split}.parquet")
        
        # Salvar o dataset como .parquet
        try:
            split_data.to_parquet(output_path)
            print(f"    Salvado em {output_path}")
        except Exception as e:
            print(f"    Erro ao salvar {output_path}: {e}")

def verify_token_ids(tokenizer, id2score):
    """
    Verifica se os token_ids mapeados correspondem aos tokens corretos.
    
    Args:
        tokenizer: O tokenizer correspondente ao modelo.
        id2score (dict): Mapeamento de token_id para pontuação.
    """
    print("Verificando mapeamento de token IDs:")
    for token_id, label in id2score.items():
        try:
            token = tokenizer.convert_ids_to_tokens(token_id)
            print(f"Token ID {token_id}: '{token}' -> {label}")
        except Exception as e:
            print(f"Erro ao converter token_id {token_id}: {e}")

def main():
    parser = argparse.ArgumentParser(description="Scorer de Qualidade para Datasets do Hugging Face")
    parser.add_argument(
        "--datasets",
        type=str,
        required=True,
        help='Lista de datasets no formato JSON. Exemplo: \'[{{ "repo": "orion-research/gsmqnaoa-pt_BR", "input": "INSTRUCTION", "output": "RESPONSE" }}, {{ "repo": "orion-research/Aura-CoT-Multilang-v1", "input": "input", "output": "output" }}]\''
    )
    parser.add_argument(
        "--output",
        type=str,
        required=True,
        help="Diretório de saída onde os datasets processados serão salvos."
    )
    parser.add_argument(
        "--cpu",
        action="store_true",
        help="Força o uso da CPU em vez da GPU."
    )
    parser.add_argument(
        "--4bit",
        action="store_true",
        dest="fourbit",
        help="Carrega o modelo em precisão de 4 bits utilizando BitsAndBytes."
    )
    
    args = parser.parse_args()
    
    # Parsear o argumento datasets
    try:
        datasets_list = json.loads(args.datasets)
        if not isinstance(datasets_list, list):
            raise ValueError("O argumento --datasets deve ser uma lista de objetos JSON.")
    except json.JSONDecodeError as e:
        print(f"Erro ao parsear o argumento --datasets: {e}")
        sys.exit(1)
    except ValueError as ve:
        print(ve)
        sys.exit(1)
    
    # Verificar se o diretório de saída existe, senão criar
    if not os.path.exists(args.output):
        try:
            os.makedirs(args.output)
            print(f"Criado diretório de saída: {args.output}")
        except Exception as e:
            print(f"Erro ao criar o diretório de saída {args.output}: {e}")
            sys.exit(1)
    
    # Determinar o dispositivo a ser usado
    if args.cpu:
        device = torch.device("cpu")
        print("Forçando o uso da CPU.")
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Usando dispositivo: {device}")
    
    # Carregar modelo e tokenizer
    model_name = "/ors/models/LLM/Orion-Quality-Scorer"
    print("Carregando tokenizer e modelo...")
    try:
        if args.fourbit:
            # Verificar se bitsandbytes está instalado
            try:
                import bitsandbytes as bnb
            except ImportError:
                print("Erro: bitsandbytes não está instalado. Instale com 'pip install bitsandbytes'.")
                sys.exit(1)
            
            from transformers import BitsAndBytesConfig
            
            # Configuração para 4-bit
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.float16
            )
            
            # Definir o device_map para permitir offloading
            device_map = "auto" if not args.cpu else {"": "cpu"}
            
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                quantization_config=bnb_config,
                device_map=device_map,
                trust_remote_code=True  # Se o modelo requer código remoto
            )
        else:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(model_name)
            model.to(device)
        
        model.eval()
        print("Modelo e tokenizer carregados com sucesso.")
        
        # Verificar mapeamento de token IDs
        verify_token_ids(tokenizer, {
            29896: "1",
            29906: "2",
            29941: "3",
            29946: "4",
            29945: "5",
            29953: "6"
        })
    except Exception as e:
        print(f"Erro ao carregar o modelo ou tokenizer: {e}")
        sys.exit(1)
    
    # Processar cada dataset
    for dataset_info in datasets_list:
        repo = dataset_info.get("repo")
        input_field = dataset_info.get("input")
        output_field = dataset_info.get("output")
        
        if not repo or not input_field or not output_field:
            print(f"Informações incompletas para o dataset: {dataset_info}. Pulando...")
            continue
        
        process_dataset(model, tokenizer, repo, input_field, output_field, args.output, device)
    
    print("\nProcessamento concluído.")

if __name__ == "__main__":
    main()