import copy import math import os import time from concurrent.futures import ThreadPoolExecutor, as_completed import requests from dotenv import load_dotenv from tqdm.auto import tqdm from .grammar_definition import apply_prompt_format, flatten # Load this project's .env when the script is launched from the project root. # Existing shell environment variables still take precedence. load_dotenv() PRINT_HIDDEN_STATE = False def call_openai_api_with_retry(args, prompt, max_tokens=10): """Call an OpenAI-compatible Chat Completions endpoint without local ML dependencies.""" url = os.getenv(args.api_url_env) api_key = os.getenv(args.api_key_env) if not url or not api_key: raise RuntimeError( f'Missing {args.api_url_env} or {args.api_key_env}. Put both values in .env or export them.') payload = { 'model': args.gpt3_engine, 'messages': [ {'role': 'system', 'content': 'You are a helpful assistant.'}, {'role': 'user', 'content': prompt}, ], 'max_tokens': max_tokens, 'temperature': 0, 'top_p': 1.0, } if ( args.api_provider == 'siliconflow' and args.gpt3_engine.startswith('Qwen/Qwen3.5-') ): payload['enable_thinking'] = False for attempt in range(4): try: response = requests.post( url, headers={'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'}, json=payload, timeout=120, ) # Only transient errors are retried. Configuration and authentication # errors should be reported immediately. if response.status_code == 429 or response.status_code >= 500: response.raise_for_status() response.raise_for_status() result = response.json() generation = result['choices'][0]['message']['content'] if not isinstance(generation, str): raise RuntimeError(f'Unexpected completion content: {generation!r}') tokens_used = result.get('usage', {}).get('total_tokens', 0) return generation.strip(), tokens_used except (requests.Timeout, requests.ConnectionError, requests.HTTPError) as error: retryable = isinstance(error, (requests.Timeout, requests.ConnectionError)) or \ getattr(error.response, 'status_code', 0) == 429 or \ getattr(error.response, 'status_code', 0) >= 500 if not retryable or attempt == 3: raise RuntimeError(f'OpenCode request failed: {error}') from error wait_seconds = 2 ** attempt print(f'OpenCode request failed ({error}); retrying in {wait_seconds}s.') time.sleep(wait_seconds) def query_model_parallelized(model, tokenizer, prompt_list, max_tokens, top_p, temperature): import torch inputs = tokenizer(prompt_list, padding=True, return_tensors='pt', return_token_type_ids=False).to('cuda') with torch.no_grad(): outputs = model.generate( **inputs, top_p=top_p, temperature=temperature, max_new_tokens=max_tokens, return_dict_in_generate=True, output_hidden_states=True, output_attentions=False, output_scores=True ) logits_list = [[] for _ in range(len(prompt_list))] # we do not print hidden state and scores because it is too much memory spenditure if PRINT_HIDDEN_STATE: # take the first (0th) inference. Its last layer (-1) will have shape [1, prompt_size, 4096]. Take last one. final_prompt_hidden_state_list = [ outputs['hidden_states'][0][-1][i, -1, :].tolist() for i in range(len(prompt_list))] else: for new_token_idx in range(len(outputs['scores'])): for i in range(len(prompt_list)): logits = torch.topk(outputs['scores'][new_token_idx][i, :], k=100) logits = [(value, index) for value, index in zip(logits.values.tolist(), logits.indices.tolist())] logits_list[i].append(logits) final_prompt_hidden_state_list = [None for _ in range(len(prompt_list))] generated_answer_list = [s.lower() for s in tokenizer.batch_decode(outputs['sequences'], skip_special_tokens=True)] return generated_answer_list, logits_list, final_prompt_hidden_state_list def _apply_prompt_format_to_extracted_fields( structured_prompt_format, input_fields_list, regex_key_idx_list, output_fields_list=None): # Precompute all format options prompt = {} for key in set(regex_key_idx_list): prompt[key] = flatten(structured_prompt_format.solve( {'enumeration_length': key, 'print_output_fields': True, 'exclude_text_field_for_output_fields': output_fields_list is None}) ).replace('<|text|>', '{}') # add empty default values if no output will be printed. It has to be a tuple to be able to concat with input_fields if output_fields_list is None: output_fields_list = [() for _ in input_fields_list] else: output_fields_list = [(output_field,) for output_field in output_fields_list] formatted_inputs = [] for input_fields, regex_key_idx, output_field in zip(input_fields_list, regex_key_idx_list, output_fields_list): tmp = apply_prompt_format(prompt[regex_key_idx], input_fields + output_field) formatted_inputs.append(tmp) return formatted_inputs def _setup_formatted_demonstrations_with_definition( structured_prompt_format, demonstration_definition, demonstrations_outputs, original_to_current_multiple_choice_classes, demos_fields_list, demos_regex_key_idx_list): # 1. replace the variables in the demonstration definition. Used when the instruction mentions # multiple choice options, which need to change when the format changes demonstration_definition = demonstration_definition.format( **structured_prompt_format.find_all_formatted_field_values() ) demonstrations_outputs = [demo[0] if isinstance(demo, list) else demo for demo in demonstrations_outputs] if original_to_current_multiple_choice_classes: demonstrations_outputs = [original_to_current_multiple_choice_classes[d] for d in demonstrations_outputs] all_demonstrations = _apply_prompt_format_to_extracted_fields( structured_prompt_format, demos_fields_list, demos_regex_key_idx_list, demonstrations_outputs) demonstration_string = demonstration_definition + "\n\n" + "\n\n".join(all_demonstrations) return demonstration_string def _setup_full_prompts_to_test_on(input_fields_list, regex_key_idx_list, selected_dataset_ids, demos_fields_list, demos_regex_key_idx_list, demonstrations_outputs, demonstration_definition, structured_prompt_format, original_to_current_multiple_choice_classes, interval_ids_to_test, n_shot): """ This function creates the full prompt string to be tested. This requires: - Formatting the demonstrations with its definition, which may require replacing some variables referring to multiple choice options. - Apply prompt format to the desired set of examples to be tested (determined by interval_ids_to_test). """ demonstration_string = _setup_formatted_demonstrations_with_definition( structured_prompt_format, demonstration_definition, demonstrations_outputs, original_to_current_multiple_choice_classes, demos_fields_list, demos_regex_key_idx_list ) # filter to keep desired interval inputs = _apply_prompt_format_to_extracted_fields( structured_prompt_format, input_fields_list[interval_ids_to_test[0]:interval_ids_to_test[1]], regex_key_idx_list[interval_ids_to_test[0]:interval_ids_to_test[1]] ) selected_dataset_ids = selected_dataset_ids[interval_ids_to_test[0]:interval_ids_to_test[1]] full_prompt_string_list = [] for input_element, idx in zip(inputs, selected_dataset_ids): full_prompt_string_list.append(input_element if n_shot == 0 else demonstration_string + "\n\n" + input_element) return full_prompt_string_list, selected_dataset_ids def evaluate_prompt_format( args, dataset, input_fields_list, regex_key_idx_list, selected_dataset_ids, demos_fields_list, demos_regex_key_idx_list, demonstrations_outputs, demonstration_definition, structured_prompt_format, model, tokenizer, model_will_repeat_input, original_to_current_multiple_choice_classes, interval_ids_to_test=(None, None)): """ Function that evaluates a prompt format (i.e. node) on a given set of samples (interval_ids_to_test). If interval_ids_to_test is not provided, it defaults to evaluating the whole dataset. """ # 1. set up input prompts including demonstrations input_prompt_string_list, selected_dataset_ids = _setup_full_prompts_to_test_on( input_fields_list, regex_key_idx_list, selected_dataset_ids, demos_fields_list, demos_regex_key_idx_list, demonstrations_outputs, demonstration_definition, structured_prompt_format, original_to_current_multiple_choice_classes, interval_ids_to_test, args.n_shot) # 2. update the output values if needed, i.e. if the multiple choice classes now have different names assert all(len(dataset[idx]['output']) == 1 for idx in selected_dataset_ids) dataset_updated = copy.deepcopy(dataset) if original_to_current_multiple_choice_classes: for idx in range(len(dataset)): dataset_updated[idx]['output'][0] = original_to_current_multiple_choice_classes[dataset[idx]['output'][0]] output_classes = sorted(list(set([dataset_updated[idx]['output'][0] for idx in selected_dataset_ids]))) # 3. evaluate if args.evaluation_metric == 'probability_ranking': return solve_with_rank_based_scoring( dataset_updated, selected_dataset_ids, model, tokenizer, input_prompt_string_list, args.batch_size_llm) elif args.evaluation_metric == 'exact_prefix_matching': logs = generate_text_with_metadata( args, input_prompt_string_list, model, tokenizer, model_will_repeat_input, dataset_updated, selected_dataset_ids, output_classes) return exact_prefix_matching_scoring(logs) def generate_text_with_metadata(args, input_prompt_string_list, model, tokenizer, model_will_repeat_input, dataset, selected_dataset_ids, output_classes): logs = [] all_tokens_used = 0 progress = tqdm( total=len(input_prompt_string_list), desc='API evaluation' if args.use_gpt3 else 'Local evaluation', unit='sample', leave=False, ) effective_batch_size = max(args.batch_size_llm, args.api_concurrency) if args.use_gpt3 else args.batch_size_llm for batch_idx in range(math.ceil(len(input_prompt_string_list) / effective_batch_size)): batch_range = [batch_idx * effective_batch_size, (batch_idx + 1) * effective_batch_size] # [) range full_prompt_string_list = input_prompt_string_list[batch_range[0]:batch_range[1]] if args.use_gpt3: request_results = [None] * len(full_prompt_string_list) with ThreadPoolExecutor(max_workers=args.api_concurrency) as executor: futures = { executor.submit(call_openai_api_with_retry, args, prompt, args.max_new_tokens): index for index, prompt in enumerate(full_prompt_string_list) } for future in as_completed(futures): index = futures[future] request_results[index] = future.result() progress.update(1) generation_list = [generation for generation, _ in request_results] all_tokens_used += sum(tokens_used for _, tokens_used in request_results) score_list = [None for _ in range(len(generation_list))] final_prompt_hidden_state_list = [None for _ in range(len(generation_list))] else: generation_list, score_list, final_prompt_hidden_state_list = query_model_parallelized( model, tokenizer, full_prompt_string_list, max_tokens=args.max_new_tokens, top_p=1.0, temperature=1.0, ) if model_will_repeat_input: generation_list = [generation[len(full_prompt_string):] for generation, full_prompt_string in zip(generation_list, full_prompt_string_list)] progress.update(len(generation_list)) selected_dataset_ids_list = [idx for idx in selected_dataset_ids[batch_range[0]:batch_range[1]]] assert len(generation_list) == len(selected_dataset_ids_list) == len(score_list) == len( final_prompt_hidden_state_list) == len(full_prompt_string_list) for generation, scores, idx, final_prompt_hidden_state, full_prompt_string in \ zip(generation_list, score_list, selected_dataset_ids_list, final_prompt_hidden_state_list, full_prompt_string_list): expected_output = dataset[idx]['output'][0] # 'entry' and 'output_classes' are needed for score generations current_log = { 'entry': dataset[idx], 'dataset_idx': idx, 'generation': generation, 'answer': expected_output, 'output_classes': output_classes, 'full_prompt_string': full_prompt_string, 'eval_type': 'exact_prefix_matching', 'scores': scores, } if PRINT_HIDDEN_STATE: current_log['final_prompt_hidden_state'] = final_prompt_hidden_state logs.append(current_log) progress.close() print('Total tokens used:', all_tokens_used) return logs def match_robust_to_multiple_choice(generation, answer_to_compare): """ We return whether the generation matched with the expected answer. This function assumes clean_text has already been run. """ # likewise, if the response says "article" and the right answer is "a" if not generation.startswith(answer_to_compare): return False # if generation starts with answer and they are the same length, they are the same string if len(generation) == len(answer_to_compare): return True # if the generation starts with the correct text, make sure the next char is not text or number # otherwise it might be just the first part of a random word (e.g. "a" with "article") # or if correct answer is ii, and all answers are i, ii, iii, iv, avoid being overly optimistic! return not generation[len(answer_to_compare)].isalpha() and not generation[len(answer_to_compare)].isdigit() def exact_prefix_matching_scoring(logs): accuracy = { 'right': [], 'wrong': [], 'other': [], 'total': 0 } for entry in logs: clean_text = lambda x: x.strip(' .,()\n-><').lower() right_answer = entry['entry']['output'][0] wrong_answers = [e for e in entry['output_classes'] if e != right_answer] entry['right_answer_formatted'] = right_answer entry['wrong_answers_formatted'] = wrong_answers right_answer = clean_text(right_answer) wrong_answers = [clean_text(e) for e in wrong_answers] generation = entry['generation'] clean_generation = clean_text(generation) is_right = match_robust_to_multiple_choice(clean_generation, right_answer) is_wrong = any( match_robust_to_multiple_choice(clean_generation, wrong_answer) for wrong_answer in wrong_answers) accuracy['right'].append(is_right) accuracy['wrong'].append(is_wrong) accuracy['other'].append(not is_wrong and not is_right) accuracy['total'] += 1 if 'output_classes' in entry and len(entry['output_classes']) > 50: del entry['output_classes'] # not changing this since it's called from many classes return (sum(accuracy['right']) * 1.0 / max(accuracy['total'], 1), sum(accuracy['wrong']) * 1.0 / max(accuracy['total'], 1), accuracy['total']), (accuracy, logs) def solve_with_rank_based_scoring( dataset, selected_dataset_ids, model, tokenizer, input_prompt_string_list, batch_size_llm): import psutil output_classes = sorted(list(set([dataset[idx]['output'][0] for idx in selected_dataset_ids]))) assert len(output_classes) < 100 assert tokenizer is not None and model is not None # if all output values are only one token, then we can just look at the output probabilities # instead of computing perplexity for all possible prompt+outputs! # also if all output values share the same prefix. E.g. ['0', '1'] tokenizes to [[1, 29871, 29900], [1, 29871, 29896]] # the first token id is always '1', so we ignore it output_classes_tokens = [t for t in tokenizer(output_classes, return_token_type_ids=False)['input_ids']] single_token_classes = all([len(t) == 2 for t in output_classes_tokens]) all_classes_share_common_prefix = len(set([tuple(t[:-1]) for t in output_classes_tokens])) == 1 accuracy = { 'right': [], 'wrong': [], 'other': [], 'total': 0 } logs = [] if single_token_classes or all_classes_share_common_prefix: # batching happens across inputs for batch_idx in range(math.ceil(len(input_prompt_string_list) / batch_size_llm)): print("Memory usage:", psutil.Process(os.getpid()).memory_info().rss / 1024 ** 2) batch_range = [batch_idx * batch_size_llm, (batch_idx + 1) * batch_size_llm] # [) range full_prompt_string_list = input_prompt_string_list[batch_range[0]:batch_range[1]] generation_list = get_ranking_based_generation_single_token_output_classes( full_prompt_string_list, output_classes, tokenizer, model) selected_dataset_ids_list = [idx for idx in selected_dataset_ids[batch_range[0]:batch_range[1]]] assert len(generation_list) == len(selected_dataset_ids_list), f"{len(generation_list)} generations, {len(selected_dataset_ids_list)} selected ids" assert len(generation_list) == len(full_prompt_string_list) for generation, idx, full_prompt_string in zip(generation_list, selected_dataset_ids_list, full_prompt_string_list): expected_output = dataset[idx]['output'][0] assert expected_output in output_classes, f"expected_output={expected_output}, output_classes={output_classes}" accuracy['right'].append((generation == expected_output)) accuracy['wrong'].append((generation != expected_output and generation in output_classes)) accuracy['other'].append((generation not in output_classes)) accuracy['total'] += 1 logs.append( { 'entry': dataset[idx], 'dataset_idx': idx, 'generation': generation, 'answer': expected_output, 'output_classes': output_classes, 'full_prompt_string': full_prompt_string, 'eval_type': 'ranking_single_token', 'scores': None, } ) else: # batching happens inside each input, since we need to do inference for each prompt+possible_output for i in range(len(input_prompt_string_list)): idx = selected_dataset_ids[i] full_prompt_string = input_prompt_string_list[i] generation = get_ranking_based_generation_multiple_token_output_classes( full_prompt_string, output_classes, tokenizer, model, batch_size_llm, ) expected_output = dataset[idx]['output'][0] assert expected_output in output_classes, f"expected_output={expected_output}, output_classes={output_classes}" accuracy['right'].append((generation == expected_output)) accuracy['wrong'].append((generation != expected_output and generation in output_classes)) accuracy['other'].append((generation not in output_classes)) accuracy['total'] += 1 logs.append( { 'entry': dataset[idx], 'dataset_idx': idx, 'generation': generation, 'answer': expected_output, 'output_classes': output_classes, 'full_prompt_string': full_prompt_string, 'eval_type': 'ranking_multiple_token', 'scores': None, } ) return (sum(accuracy['right']) * 1.0 / max(accuracy['total'], 1), sum(accuracy['wrong']) * 1.0 / max(accuracy['total'], 1), accuracy['total']), (accuracy, logs) def get_ranking_based_generation_single_token_output_classes(prompts, output_classes, tokenizer, model): import torch top_p = 1.0 temperature = 1.0 # if all output values are only one token, then we can just look at the output probabilities! # also if all output values share the same prefix. E.g. ['0', '1'] tokenizes to [[1, 29871, 29900], [1, 29871, 29896]] # the first token id is always '1', so we ignore it output_classes_tokens = [t for t in tokenizer(output_classes, return_token_type_ids=False)['input_ids']] all_classes_share_common_prefix = len(set([tuple(t[:-1]) for t in output_classes_tokens])) == 1 tokenized_inputs_list = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False)[ 'input_ids'].tolist() if all_classes_share_common_prefix: for i in range(len(tokenized_inputs_list)): # if the tokenized element is [1, 29871, 29900], get [29871] tokenized_inputs_list[i] += output_classes_tokens[0][1:-1] tokenized_inputs = torch.tensor(tokenized_inputs_list).to('cuda') with torch.no_grad(): outputs = model.generate(input_ids=tokenized_inputs, top_p=top_p, temperature=temperature, max_new_tokens=1, return_dict_in_generate=True, output_scores=True) scores = outputs["scores"][0] # first dimension = 1 since we only generate one token generations = [] for i in range(len(prompts)): all_logits = scores[i, :].squeeze().tolist() all_logits_sorted = sorted([(all_logits[t[-1]], i) for i, t in enumerate(output_classes_tokens)], reverse=True) generations.append(output_classes[all_logits_sorted[0][1]]) return generations def get_ranking_based_generation_multiple_token_output_classes(prompt, output_classes, tokenizer, model, batch_size_llm): import torch import torch.nn.functional as F output_classes_tokens = [t for t in tokenizer(output_classes, return_token_type_ids=False)['input_ids']] prompts = [prompt + class_seq for class_seq in output_classes] all_logits_list, all_tokens_list = [], [] for batch_idx in range(math.ceil(len(prompts) / batch_size_llm)): batch_range = [batch_idx * batch_size_llm, (batch_idx + 1) * batch_size_llm] # [) range all_logits, all_tokens = _get_input_logits_and_tokens(prompts[batch_range[0]:batch_range[1]], tokenizer, model) all_logits_list.extend(all_logits) all_tokens_list.extend(all_tokens) n_classes = len(output_classes) class_logprobs = [] for class_index in range(n_classes): class_logits = all_logits_list[class_index] # the lengths of each class sequence in tokens target_token_length = (len(output_classes_tokens[class_index])) # we only need the logits for the end sequence tokens = all_tokens_list[class_index] # we have to go back by one because we don't care about the logits for the predicted token sequence_logits = class_logits[-target_token_length - 1: -1] sequence_tokens = tokens[-target_token_length:] # we take a log_softmax over all token logits for each position in the class sequence to # get log probabilities, and then sum the logprobs for the tokens actually chosen logprobs = F.log_softmax(sequence_logits, dim=-1).to('cpu') class_logprob = sum( [logprobs[i, token] for i, token in enumerate(sequence_tokens)] ) class_logprobs.append(class_logprob.item()) return output_classes[torch.tensor(class_logprobs).argmax(dim=-1).item()] def _get_input_logits_and_tokens(inputs, tokenizer, model): import torch tokenized_inputs = tokenizer(inputs, return_tensors="pt", padding=True, return_token_type_ids=False).to('cuda') with torch.no_grad(): outputs = model(**tokenized_inputs) logits = outputs["logits"].detach().to(device="cpu", dtype=torch.float32) return logits, tokenized_inputs["input_ids"]