mirror of
https://github.com/deepseek-ai/DeepSeek-Math
synced 2025-06-26 18:16:20 +00:00
init
This commit is contained in:
ZhihongShao
263
evaluation/eval/utils.py
Executable file
263
evaluation/eval/utils.py
Executable file
@@ -0,0 +1,263 @@
|
||||
import torch
|
||||
import tqdm
|
||||
from transformers import StoppingCriteria, GenerationConfig
|
||||
|
||||
class KeyWordsCriteria(StoppingCriteria):
|
||||
def __init__(self, stop_id_sequences, tokenizer, prompt_length):
|
||||
assert isinstance(stop_id_sequences[0], list), "stop_id_sequences should be a list of list of ids"
|
||||
self.tokenizer = tokenizer
|
||||
self.stop_id_sequences = stop_id_sequences
|
||||
self.stop_sequences = [tokenizer.decode(sequence) for sequence in stop_id_sequences]
|
||||
print(f"stop sequences: {self.stop_sequences}", flush=True)
|
||||
self.prompt_length = prompt_length
|
||||
|
||||
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
|
||||
sequences_should_be_stopped = []
|
||||
for i in range(input_ids.shape[0]):
|
||||
ids = input_ids[i][self.prompt_length:].tolist()
|
||||
should_be_stopped = False
|
||||
for stop_ids, stop_sequence in zip(self.stop_id_sequences, self.stop_sequences):
|
||||
_ids = ids
|
||||
for j in range(len(_ids), 0, -1):
|
||||
s = self.tokenizer.decode(_ids[max(j - len(stop_ids) - 3, 0) :j])
|
||||
if s.endswith(stop_sequence):
|
||||
should_be_stopped = True
|
||||
break
|
||||
if should_be_stopped:
|
||||
break
|
||||
sequences_should_be_stopped.append(should_be_stopped)
|
||||
return all(sequences_should_be_stopped)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_completions(model, tokenizer, prompts, batch_size=1, stop_id_sequences=None, end_of_generation_id_sequence=None, disable_tqdm=False, **generation_kwargs):
|
||||
generations = []
|
||||
finish_completion = []
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(prompts), desc="Generating Completions")
|
||||
|
||||
if stop_id_sequences is not None:
|
||||
stop_sequences = [tokenizer.decode(stop_id_sequence) for stop_id_sequence in stop_id_sequences]
|
||||
|
||||
if end_of_generation_id_sequence is not None:
|
||||
end_of_generation_sequence = tokenizer.decode(end_of_generation_id_sequence)
|
||||
|
||||
num_return_sequences = generation_kwargs.get("num_return_sequences", 1)
|
||||
generation_kwargs['use_cache'] = True
|
||||
for i in range(0, len(prompts), batch_size):
|
||||
batch_prompts = prompts[i:i+batch_size]
|
||||
tokenized_prompts = tokenizer(batch_prompts, padding="longest", return_tensors="pt", add_special_tokens='chatglm2' in str(model.__class__))
|
||||
batch_input_ids = tokenized_prompts.input_ids
|
||||
attention_mask = tokenized_prompts.attention_mask
|
||||
|
||||
if model.device.type == "cuda":
|
||||
batch_input_ids = batch_input_ids.cuda()
|
||||
attention_mask = attention_mask.cuda()
|
||||
|
||||
batch_finish_completion = [False] * len(batch_prompts) * num_return_sequences
|
||||
try:
|
||||
batch_outputs = model.generate(
|
||||
input_ids=batch_input_ids,
|
||||
attention_mask=attention_mask,
|
||||
stopping_criteria=[KeyWordsCriteria(stop_id_sequences, tokenizer, batch_input_ids.size(1))] if stop_id_sequences else None,
|
||||
**generation_kwargs
|
||||
)
|
||||
|
||||
# the stopping criteria is applied at batch level, so if other examples are not stopped, the entire batch will continue to generate.
|
||||
# so some outputs still have the stop sequence, which we need to remove.
|
||||
if stop_id_sequences:
|
||||
for output_idx in range(batch_outputs.shape[0]):
|
||||
finish = False
|
||||
for token_idx in range(batch_input_ids.shape[1], batch_outputs.shape[1]):
|
||||
if any(tokenizer.decode(batch_outputs[output_idx, token_idx: token_idx + len(stop_sequence) + 3]).startswith(stop_sequence) for stop_sequence in stop_sequences):
|
||||
if end_of_generation_id_sequence is not None and tokenizer.decode(batch_outputs[output_idx, token_idx: token_idx + len(end_of_generation_id_sequence) + 3]).startswith(end_of_generation_sequence):
|
||||
batch_finish_completion[output_idx] = True
|
||||
batch_outputs[output_idx, token_idx:] = tokenizer.pad_token_id
|
||||
break
|
||||
|
||||
# remove the prompt from the output
|
||||
# we need to re-encode the prompt because we need to make sure the special tokens are treated the same way as in the outputs.
|
||||
# we changed our previous way of truncating the output token ids dicrectly because some tokenizer (e.g., llama) won't add space token before the first token.
|
||||
# space is important for some tasks (e.g., code completion).
|
||||
batch_outputs = tokenizer.batch_decode(batch_outputs, skip_special_tokens=True)
|
||||
batch_prompts = tokenizer.batch_decode(batch_input_ids, skip_special_tokens=True)
|
||||
# duplicate the prompts to match the number of return sequences
|
||||
batch_prompts = [prompt for prompt in batch_prompts for _ in range(num_return_sequences)]
|
||||
batch_generations = [
|
||||
output[len(prompt):] for prompt, output in zip(batch_prompts, batch_outputs)
|
||||
]
|
||||
except Exception as e:
|
||||
print("Error when generating completions for batch:")
|
||||
print(batch_prompts)
|
||||
print("Error message:")
|
||||
print(e)
|
||||
print("Use empty string as the completion.")
|
||||
batch_generations = [""] * len(batch_prompts) * num_return_sequences
|
||||
|
||||
generations += batch_generations
|
||||
finish_completion += batch_finish_completion
|
||||
|
||||
if not disable_tqdm:
|
||||
progress.update(len(batch_prompts)//num_return_sequences)
|
||||
|
||||
assert len(generations) == len(prompts) * num_return_sequences, "number of generations should be equal to number of prompts * num_return_sequences"
|
||||
return generations, finish_completion
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def get_next_word_predictions(model, tokenizer, prompts, candidate_token_ids=None, batch_size=1, return_token_predictions=False, disable_tqdm=False):
|
||||
predictions, probs = [], []
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(prompts), desc="Getting Predictions")
|
||||
|
||||
for i in range(0, len(prompts), batch_size):
|
||||
batch_prompts = prompts[i: i+batch_size]
|
||||
tokenized_prompts = tokenizer(batch_prompts, padding="longest", return_tensors="pt", add_special_tokens=False)
|
||||
batch_input_ids = tokenized_prompts.input_ids
|
||||
attention_mask = tokenized_prompts.attention_mask
|
||||
|
||||
if model.device.type == "cuda":
|
||||
batch_input_ids = batch_input_ids.cuda()
|
||||
attention_mask = attention_mask.cuda()
|
||||
|
||||
batch_logits = model(input_ids=batch_input_ids, attention_mask=attention_mask).logits[:, -1, :]
|
||||
if candidate_token_ids is not None:
|
||||
batch_logits = batch_logits[:, candidate_token_ids]
|
||||
batch_probs = torch.softmax(batch_logits, dim=-1)
|
||||
batch_prediction_indices = torch.argmax(batch_probs, dim=-1)
|
||||
if return_token_predictions:
|
||||
if candidate_token_ids is not None:
|
||||
candidate_tokens = tokenizer.convert_ids_to_tokens(candidate_token_ids)
|
||||
batch_predictions = [candidate_tokens[idx] for idx in batch_prediction_indices]
|
||||
else:
|
||||
batch_predictions = tokenizer.convert_ids_to_tokens(batch_prediction_indices)
|
||||
predictions += batch_predictions
|
||||
else:
|
||||
predictions += batch_prediction_indices.tolist()
|
||||
probs += batch_probs.tolist()
|
||||
|
||||
if not disable_tqdm:
|
||||
progress.update(len(batch_prompts))
|
||||
|
||||
assert len(predictions) == len(prompts), "number of predictions should be equal to number of prompts"
|
||||
return predictions, probs
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def score_completions(model, tokenizer, scoring_examples, disable_tqdm=False):
|
||||
'''
|
||||
Each scoring example is a dict, which contains the following keys:
|
||||
- prompt: the prompt to score
|
||||
- completions: a list of completions to score
|
||||
'''
|
||||
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(scoring_examples), desc="Scoring Completions")
|
||||
|
||||
# unroll the scoring examples
|
||||
unrolled_examples = []
|
||||
for scoring_example in scoring_examples:
|
||||
prompt = scoring_example["prompt"]
|
||||
for completion in scoring_example["completions"]:
|
||||
unrolled_examples.append({
|
||||
"prompt": prompt,
|
||||
"completion": completion
|
||||
})
|
||||
|
||||
scores = []
|
||||
# currently we don't support batching, because we want to directly use the loss returned by the model to score each completion.
|
||||
for unrolled_example in unrolled_examples:
|
||||
encoded_example = encode_with_prompt_completion_format(unrolled_example, tokenizer, max_seq_length=None)
|
||||
# unsqueeze the batch dimension
|
||||
for key, value in encoded_example.items():
|
||||
encoded_example[key] = value.unsqueeze(0)
|
||||
if model.device.type == "cuda":
|
||||
encoded_example = {
|
||||
key: value.cuda() for key, value in encoded_example.items()
|
||||
}
|
||||
outputs = model(**encoded_example)
|
||||
loss = outputs.loss
|
||||
scores.append(-loss.item())
|
||||
if not disable_tqdm:
|
||||
progress.update(1)
|
||||
|
||||
# roll up the scores
|
||||
rolled_up_scores = {}
|
||||
for unrolled_example, score in zip(unrolled_examples, scores):
|
||||
prompt = unrolled_example["prompt"]
|
||||
completion = unrolled_example["completion"]
|
||||
if prompt not in rolled_up_scores:
|
||||
rolled_up_scores[prompt] = {}
|
||||
rolled_up_scores[prompt][completion] = score
|
||||
|
||||
return rolled_up_scores
|
||||
|
||||
|
||||
|
||||
def load_hf_lm_and_tokenizer(
|
||||
model_name_or_path,
|
||||
tokenizer_name_or_path=None,
|
||||
device_map="auto",
|
||||
load_in_8bit=False,
|
||||
load_in_half=False,
|
||||
gptq_model=False,
|
||||
use_fast_tokenizer=True,
|
||||
padding_side="left",
|
||||
):
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoModel, AutoTokenizer
|
||||
|
||||
if not tokenizer_name_or_path:
|
||||
tokenizer_name_or_path = model_name_or_path
|
||||
|
||||
is_chatglm2 = 'chatglm2' in tokenizer_name_or_path.lower() or 'chatglm2' in model_name_or_path
|
||||
is_qwen = 'qwen' in tokenizer_name_or_path.lower() or 'qwen' in model_name_or_path
|
||||
|
||||
if is_chatglm2 or is_qwen:
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True)
|
||||
if is_qwen:
|
||||
tokenizer.eos_token = '<|endoftext|>'
|
||||
tokenizer.eos_token_id = 151643
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
else:
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True, use_fast=use_fast_tokenizer)
|
||||
# set padding side to left for batch generation
|
||||
tokenizer.padding_side = padding_side
|
||||
# set pad token to eos token if pad token is not set (as is the case for llama models)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
|
||||
if gptq_model:
|
||||
from auto_gptq import AutoGPTQForCausalLM
|
||||
model_wrapper = AutoGPTQForCausalLM.from_quantized(
|
||||
model_name_or_path, device="cuda:0", use_triton=True
|
||||
)
|
||||
model = model_wrapper.model
|
||||
elif load_in_8bit:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name_or_path,
|
||||
device_map=device_map,
|
||||
load_in_8bit=True
|
||||
)
|
||||
else:
|
||||
kwargs = {}
|
||||
model_class = AutoModelForCausalLM
|
||||
if is_chatglm2:
|
||||
kwargs = {'trust_remote_code': True}
|
||||
model_class = AutoModel
|
||||
elif is_qwen:
|
||||
kwargs = {'trust_remote_code': True}
|
||||
if device_map:
|
||||
model = model_class.from_pretrained(model_name_or_path, device_map=device_map, **kwargs)
|
||||
else:
|
||||
model = model_class.from_pretrained(model_name_or_path, **kwargs)
|
||||
if torch.cuda.is_available():
|
||||
model = model.cuda()
|
||||
if is_qwen:
|
||||
model.generation_config = GenerationConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
|
||||
model.generation_config.do_sample = False
|
||||
if not is_chatglm2 and not is_qwen and load_in_half:
|
||||
model = model.half()
|
||||
model.eval()
|
||||
return model, tokenizer
|
||||
Reference in New Issue
Block a user