update tokenize function

2 weeks ago · b6ec337f3d
parent 794e0d4f4a
commit b6ec337f3d
1 changed files with 2 additions and 2 deletions
--- a/applications/ColossalChat/coati/dataset/tokenization_utils.py
+++ b/applications/ColossalChat/coati/dataset/tokenization_utils.py
@ -431,8 +431,8 @@ def tokenize_process_reward(
    decoded_label = tokenizer.decode(label, skip_special_tokens=False)

    return {
-        "input": tokenized,
-        "label": label,
+        "input_ids": tokenized,
+        "labels": label,
        "loss_mask": loss_mask,
        "decoded_input": decoded_input,
        "decoded_label": decoded_label,