511 lines
25 KiB
Python
511 lines
25 KiB
Python
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"]
|