From ca38148acdac56467636072b29d123cb7dd05231 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 24 Dec 2024 10:06:41 +0000 Subject: [PATCH] fix generation decoding --- hyperlora/intx_sft.py | 26 +++++++++++++++----------- hyperlora/training_utils.py | 30 +++++++++++++++++------------- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index f66523e..c2f4b48 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -218,6 +218,7 @@ def main(output_dir: str): add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel) # for sft + chat_model, we need to convert the dataset to chat format # add "messages" field + # TODO: apply different conversion for generation? ds = ds.map( convert_ctx_prompt_response_to_messages, fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat}, @@ -276,7 +277,7 @@ def main(output_dir: str): pad_to_multiple_of=8, return_tensors="pt", ) - labels = [x.pop("labels") for x in inp_list] + ctx_features = None if "ctx_features" in inp_list[0]: # have to be manual since it has [ctx_len, features] shape @@ -300,6 +301,7 @@ def main(output_dir: str): padded_seq = tokenizer.pad(inp_list, **padding_kwargs) # hacky explicit padding since the labels are not padded by default + labels = [x.pop("labels") for x in inp_list] labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"] labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels) out = {**padded_seq, "labels": labels} @@ -308,28 +310,30 @@ def main(output_dir: str): out["ctx_attn_mask"] = ctx_attn_mask return out + # TODO: generation dataset shouldn't include labels in the "input_ids" field def generation_collator(inp_list, tokenizer): padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt") input_ids = [x.pop("input_ids") for x in inp_list] attn_mask = [x.pop("attention_mask") for x in inp_list] labels = [x.pop("labels") for x in inp_list] for i, label in enumerate(labels): - # remove the label part + # HACK: remove the label part idx = np.argmax(label != -100) input_ids[i] = input_ids[i][:idx] attn_mask[i] = attn_mask[i][:idx] out = tokenizer.pad( {"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs ) - label_pad_len = len(out["input_ids"][0]) - labels[0] = torch.cat( - [torch.tensor([-100] * (label_pad_len - len(label))), label], dim=0 - ) - labels = torch.nn.utils.rnn.pad_sequence( - labels, - batch_first=True, - padding_value=-100, - ).long() + # label_pad_len = len(out["input_ids"][0]) + # labels[0] = torch.cat( + # [torch.tensor([-100] * (label_pad_len - len(label))), label], dim=0 + # ) + # labels = torch.nn.utils.rnn.pad_sequence( + # labels, + # batch_first=True, + # padding_value=-100, + # ).long() + out["labels"] = labels if "ctx_features" in inp_list[0]: diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py index 0f82499..e5c2536 100644 --- a/hyperlora/training_utils.py +++ b/hyperlora/training_utils.py @@ -24,20 +24,24 @@ def save_generated_text(samples, output_dir, split): def decode_test_result(test_dataset, test_result, tokenizer): - for sample, pred_toks, labels in zip( - test_dataset, test_result.predictions, test_result.label_ids - ): - start_idx = np.argmax(labels != -100, axis=0) - input_toks = sample["input_ids"][:start_idx] - gen_toks = pred_toks[start_idx:] - label_toks = labels[start_idx:] - # labels are padded with -100, so we need to replace them with the pad token id - label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks) + out = dict() + for sample, pred_toks in zip(test_dataset, test_result.predictions): + if "labels" in sample: + start_idx = np.argmax(sample["labels"] != -100) + label_toks = sample["labels"][start_idx:] + # labels are padded with -100, so we need to replace them with the pad token id + label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks) + label_text = tokenizer.decode(label_toks, skip_special_tokens=True) + out["label"] = label_text - input_text = tokenizer.decode(input_toks, skip_special_tokens=True) - gen_text = tokenizer.decode(gen_toks, skip_special_tokens=True) - label_text = tokenizer.decode(label_toks, skip_special_tokens=True) - yield {"input": input_text, "generated": gen_text, "label": label_text} + # HACK: remove the label part + input_toks = sample["input_ids"][:start_idx] + gen_toks = pred_toks[len(input_toks) :] + + out["input"] = tokenizer.decode(input_toks, skip_special_tokens=True) + out["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True) + + yield out def eval_generation(eval_trainer, tokenizer, dataset, split, gen_kwargs):