Views
No views yet
1#!/usr/bin/env python3
2import sys
3import os
4
5# Passo 1: Parsear os argumentos de linha de comando antecipadamente para definir variáveis de ambiente
6# Antes de importar transformers, parseamos os argumentos que podem afetar a configuração do modelo
7def early_parse_args():
8 import argparse
9
10 parser = argparse.ArgumentParser(description="Scorer de Qualidade para Datasets do Hugging Face (Early Parsing)")
11 parser.add_argument(
12 "--4bit",
13 action="store_true",
14 dest="fourbit",
15 help="Carrega o modelo em precisão de 4 bits utilizando BitsAndBytes."
16 )
17 # Parse apenas os argumentos necessários para BitsAndBytes
18 args, _ = parser.parse_known_args()
19 return args
20
21early_args = early_parse_args()
22
23# Definir a variável de ambiente antes de importar transformers, se --4bit for usado
24if early_args.fourbit:
25 os.environ["LLM_INT8_ENABLE_FP32_CPU_OFFLOAD"] = "true"
26
27# Passo 2: Importar as bibliotecas restantes após definir variáveis de ambiente
28import argparse
29import json
30from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
31import numpy as np
32from scipy.special import softmax
33from datasets import load_dataset
34import pyarrow.parquet as pq
35from tqdm import tqdm
36import torch
37
38def infer_quality(model, tokenizer, input_text, resp_text, device):
39 """
40 Calcula a pontuação de qualidade para uma resposta baseada na entrada e na resposta fornecida.
41
42 Args:
43 model: O modelo de linguagem causal carregado.
44 tokenizer: O tokenizer correspondente ao modelo.
45 input_text (str): O texto da instrução/pergunta.
46 resp_text (str): O texto da resposta a ser avaliada.
47 device: O dispositivo (CPU ou GPU) para realizar os cálculos.
48
49 Returns:
50 float: A pontuação de qualidade calculada.
51 """
52 # Template traduzido para português
53 quality_template = (
54 "Você é um assistente útil. Por favor, identifique a pontuação de qualidade da Resposta correspondente à Pergunta.\n"
55 "#Pergunta#:\n{instruction}\n"
56 "#Resposta#:\n{output}\n"
57 "##Qualidade: "
58 )
59
60 # Formatar o input do usuário
61 user_input = quality_template.format(instruction=input_text, output=resp_text)
62
63 # Tokenizar a entrada
64 input_ids = tokenizer.encode(user_input, return_tensors="pt").to(device)
65 max_length = 512 # Definir comprimento máximo
66
67 # Gerar a resposta do modelo
68 with torch.no_grad():
69 outputs = model.generate(
70 input_ids,
71 max_length=max_length,
72 num_return_sequences=1,
73 return_dict_in_generate=True,
74 output_scores=True,
75 do_sample=False # Desativar amostragem para obter logits determinísticos
76 )
77
78 # Extrair os scores das saídas
79 scores = outputs.scores # Tuple of tensors, one for each generated token
80 score_logits = []
81
82 # Mapeamento de IDs de tokens para pontuações
83 id2score = {
84 29896: "1", # Token ID para "1"
85 29906: "2", # Token ID para "2"
86 29941: "3", # Token ID para "3"
87 29946: "4", # Token ID para "4"
88 29945: "5", # Token ID para "5"
89 29953: "6" # Token ID para "6"
90 }
91
92 score_template = np.array([1, 2, 3, 4, 5, 6])
93
94 # Verificar se há scores gerados
95 if not scores:
96 return 0.0 # Retorna 0 se não houver scores
97
98 # Assumindo que queremos a última posição gerada
99 last_score = scores[-1] # Último conjunto de logits, shape: (batch_size * num_beams, vocab_size)
100
101 # Garantir que last_score tenha o formato esperado
102 if last_score.ndim != 2 or last_score.size(0) != 1:
103 print("Formato inesperado para 'last_score'.")
104 return 0.0
105
106 last_score = last_score[0] # Shape: (vocab_size,)
107
108 # Extrair logits para os token_ids específicos
109 for token_id in id2score:
110 if token_id < last_score.size(0):
111 logit = last_score[token_id].item()
112 score_logits.append(logit)
113 else:
114 # Se o token_id estiver fora do alcance, atribuir um valor muito baixo
115 score_logits.append(-1e10)
116
117 score_logits = np.array(score_logits)
118
119 # Aplicar softmax para obter probabilidades
120 score_probs = softmax(score_logits)
121
122 # Calcular a pontuação ponderada
123 quality_score = np.sum(score_probs * score_template)
124
125 return quality_score
126
127def process_dataset(model, tokenizer, dataset_repo, input_field, output_field, output_dir, device):
128 """
129 Processa um dataset específico, calculando a pontuação de qualidade para cada exemplo e salvando o resultado.
130
131 Args:
132 model: O modelo de linguagem causal carregado.
133 tokenizer: O tokenizer correspondente ao modelo.
134 dataset_repo (str): O repositório do dataset no Hugging Face.
135 input_field (str): O nome do campo de entrada no dataset.
136 output_field (str): O nome do campo de resposta no dataset.
137 output_dir (str): O diretório onde os datasets processados serão salvos.
138 device: O dispositivo (CPU ou GPU) para realizar os cálculos.
139 """
140 print(f"\nProcessando dataset: {dataset_repo}")
141 try:
142 dataset = load_dataset(dataset_repo)
143 except Exception as e:
144 print(f"Erro ao carregar o dataset {dataset_repo}: {e}")
145 return
146
147 # Considerar apenas splits que contêm dados (por exemplo, 'train', 'validation', 'test')
148 for split in dataset.keys():
149 split_data = dataset[split]
150 if len(split_data) == 0:
151 continue
152
153 print(f" Processando split: {split} com {len(split_data)} exemplos")
154
155 # Preparar listas para armazenar os scores
156 quality_scores = []
157
158 # Iterar sobre os exemplos com tqdm para visualizar o progresso
159 for example in tqdm(split_data, desc=f" Avaliando {split}"):
160 input_text = example.get(input_field, "")
161 output_text = example.get(output_field, "")
162 if not isinstance(input_text, str) or not isinstance(output_text, str):
163 quality_scores.append(0.0)
164 continue
165 score = infer_quality(model, tokenizer, input_text, output_text, device)
166 quality_scores.append(score)
167
168 # Adicionar a nova coluna ao dataset
169 split_data = split_data.add_column("quality_score", quality_scores)
170
171 # Definir o caminho de salvamento
172 dataset_name = dataset_repo.split('/')[-1]
173 split_output_dir = os.path.join(output_dir, dataset_name)
174 os.makedirs(split_output_dir, exist_ok=True)
175 output_path = os.path.join(split_output_dir, f"{split}.parquet")
176
177 # Salvar o dataset como .parquet
178 try:
179 split_data.to_parquet(output_path)
180 print(f" Salvado em {output_path}")
181 except Exception as e:
182 print(f" Erro ao salvar {output_path}: {e}")
183
184def verify_token_ids(tokenizer, id2score):
185 """
186 Verifica se os token_ids mapeados correspondem aos tokens corretos.
187
188 Args:
189 tokenizer: O tokenizer correspondente ao modelo.
190 id2score (dict): Mapeamento de token_id para pontuação.
191 """
192 print("Verificando mapeamento de token IDs:")
193 for token_id, label in id2score.items():
194 try:
195 token = tokenizer.convert_ids_to_tokens(token_id)
196 print(f"Token ID {token_id}: '{token}' -> {label}")
197 except Exception as e:
198 print(f"Erro ao converter token_id {token_id}: {e}")
199
200def main():
201 parser = argparse.ArgumentParser(description="Scorer de Qualidade para Datasets do Hugging Face")
202 parser.add_argument(
203 "--datasets",
204 type=str,
205 required=True,
206 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" }}]\''
207 )
208 parser.add_argument(
209 "--output",
210 type=str,
211 required=True,
212 help="Diretório de saída onde os datasets processados serão salvos."
213 )
214 parser.add_argument(
215 "--cpu",
216 action="store_true",
217 help="Força o uso da CPU em vez da GPU."
218 )
219 parser.add_argument(
220 "--4bit",
221 action="store_true",
222 dest="fourbit",
223 help="Carrega o modelo em precisão de 4 bits utilizando BitsAndBytes."
224 )
225
226 args = parser.parse_args()
227
228 # Parsear o argumento datasets
229 try:
230 datasets_list = json.loads(args.datasets)
231 if not isinstance(datasets_list, list):
232 raise ValueError("O argumento --datasets deve ser uma lista de objetos JSON.")
233 except json.JSONDecodeError as e:
234 print(f"Erro ao parsear o argumento --datasets: {e}")
235 sys.exit(1)
236 except ValueError as ve:
237 print(ve)
238 sys.exit(1)
239
240 # Verificar se o diretório de saída existe, senão criar
241 if not os.path.exists(args.output):
242 try:
243 os.makedirs(args.output)
244 print(f"Criado diretório de saída: {args.output}")
245 except Exception as e:
246 print(f"Erro ao criar o diretório de saída {args.output}: {e}")
247 sys.exit(1)
248
249 # Determinar o dispositivo a ser usado
250 if args.cpu:
251 device = torch.device("cpu")
252 print("Forçando o uso da CPU.")
253 else:
254 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
255 print(f"Usando dispositivo: {device}")
256
257 # Carregar modelo e tokenizer
258 model_name = "/ors/models/LLM/Orion-Quality-Scorer"
259 print("Carregando tokenizer e modelo...")
260 try:
261 if args.fourbit:
262 # Verificar se bitsandbytes está instalado
263 try:
264 import bitsandbytes as bnb
265 except ImportError:
266 print("Erro: bitsandbytes não está instalado. Instale com 'pip install bitsandbytes'.")
267 sys.exit(1)
268
269 from transformers import BitsAndBytesConfig
270
271 # Configuração para 4-bit
272 bnb_config = BitsAndBytesConfig(
273 load_in_4bit=True,
274 bnb_4bit_use_double_quant=True,
275 bnb_4bit_quant_type="nf4",
276 bnb_4bit_compute_dtype=torch.float16
277 )
278
279 # Definir o device_map para permitir offloading
280 device_map = "auto" if not args.cpu else {"": "cpu"}
281
282 tokenizer = AutoTokenizer.from_pretrained(model_name)
283 model = AutoModelForCausalLM.from_pretrained(
284 model_name,
285 quantization_config=bnb_config,
286 device_map=device_map,
287 trust_remote_code=True # Se o modelo requer código remoto
288 )
289 else:
290 tokenizer = AutoTokenizer.from_pretrained(model_name)
291 model = AutoModelForCausalLM.from_pretrained(model_name)
292 model.to(device)
293
294 model.eval()
295 print("Modelo e tokenizer carregados com sucesso.")
296
297 # Verificar mapeamento de token IDs
298 verify_token_ids(tokenizer, {
299 29896: "1",
300 29906: "2",
301 29941: "3",
302 29946: "4",
303 29945: "5",
304 29953: "6"
305 })
306 except Exception as e:
307 print(f"Erro ao carregar o modelo ou tokenizer: {e}")
308 sys.exit(1)
309
310 # Processar cada dataset
311 for dataset_info in datasets_list:
312 repo = dataset_info.get("repo")
313 input_field = dataset_info.get("input")
314 output_field = dataset_info.get("output")
315
316 if not repo or not input_field or not output_field:
317 print(f"Informações incompletas para o dataset: {dataset_info}. Pulando...")
318 continue
319
320 process_dataset(model, tokenizer, repo, input_field, output_field, args.output, device)
321
322 print("\nProcessamento concluído.")
323
324if __name__ == "__main__":
325 main()