#!/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()