fix generation decoding

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]:

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):