fix(offline_pipeline): ILQL negative indexing under truncation (#435)

* fix(offline_pipeline): prepend `is_output=False` msg when truncated

* fix(test_pipelines): specify `truncation_side`

* style

* fix(offline_pipeline): prepend starting <bos> under truncation

* fix(test_pipelines): update tests for the truncation change

* docs(offline_pipeline): update `tokenize_dialogue` type signature

Author: maxreciprocate

tests/test_pipelines.py

@@ -11,6 +11,19 @@ class TestTokenizeDialog(TestCase):
     def setUp(self):
         self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
+    def test_tokenize_dialogue_truncation_basic(self):
+        dialogue = ["this will be truncated", "."]
+        self.tokenizer.truncation_side = "left"
+
+        dialog = tokenize_dialogue(dialogue, self.tokenizer, max_length=2)
+
+        assert len(dialog) == 2
+        user_dm, bot_dm = dialog
+        assert len(user_dm.tokens) == 1
+        assert len(bot_dm.tokens) == 1
+        assert user_dm == DialogMessage(is_output=False, tokens=(self.tokenizer.bos_token_id,))
+        assert bot_dm == DialogMessage(is_output=True, tokens=(self.tokenizer.eos_token_id,))
+
     @given(st.lists(st.text(), max_size=32))
     def test_tokenize_dialogue_single_turn(self, response_words):
         response = " ".join(response_words)  # space seperate to make it multiple tokens
@@ -46,20 +59,18 @@ def test_tokenize_dialogue_single_turn_truncation_left(self, response_words, max
         response = " ".join(response_words)  # space seperate to make it multiple tokens
         self.tokenizer.truncation_side = "left"
         tokenized_response = tuple(self.tokenizer(response, add_special_tokens=False).input_ids)
-        tokenized_response = tokenized_response + (self.tokenizer.eos_token_id,)
+        tokenized_response += (self.tokenizer.eos_token_id,)
         dialog = tokenize_dialogue(response, self.tokenizer, max_length=max_length)
 
-        # if no truncation should have happened, then the user BOS prompt should be present
-        if len(tokenized_response) + 1 <= max_length:
-            assert len(dialog) == 2
-            user_dm, bot_dm = dialog
+        # whether or not truncation has happened, user BOS prompt should be present
+        assert len(dialog) == 2
+        user_dm, bot_dm = dialog
+        assert user_dm == DialogMessage(is_output=False, tokens=(self.tokenizer.bos_token_id,))
 
-            assert user_dm == DialogMessage(is_output=False, tokens=(self.tokenizer.bos_token_id,))
+        if len(tokenized_response) < max_length:
             assert bot_dm == DialogMessage(is_output=True, tokens=tokenized_response)
         else:
-            assert len(dialog) == 1
-            bot_dm = dialog[0]
-            assert bot_dm == DialogMessage(is_output=True, tokens=tokenized_response[-max_length:])
+            assert bot_dm == DialogMessage(is_output=True, tokens=tokenized_response[-max_length + 1 :])
 
         all_tokens = sum((dm.tokens for dm in dialog), ())
         assert len(all_tokens) <= max_length
@@ -76,6 +87,9 @@ def test_tokenize_dialogue_multi_turn(self, user_response_pairs):
 
         dm_convo = [DialogMessage(is_output=i % 2 == 1, tokens=tokens) for i, tokens in enumerate(tokenized_flat_convo)]
         nonempty_dm_convo = [dm for dm in dm_convo if dm.tokens]
+        if nonempty_dm_convo[0].is_output:
+            nonempty_dm_convo.insert(0, DialogMessage(is_output=False, tokens=(self.tokenizer.eos_token_id,)))
+
         assert dialog == nonempty_dm_convo
 
     @given(st.lists(st.tuples(st.text(), st.text()), min_size=1, max_size=32), st.integers(min_value=2, max_value=16))
@@ -91,6 +105,9 @@ def test_tokenize_dialogue_multi_turn_truncation_right(self, user_response_pairs
 
         all_tokens = sum((dm.tokens for dm in dialog), ())
         should_be_tokens = sum(tokenized_flat_convo, ())[:max_length]
+        if dialog[0] == DialogMessage(is_output=False, tokens=(self.tokenizer.eos_token_id,)):
+            should_be_tokens = (self.tokenizer.eos_token_id, *should_be_tokens[: max_length - 1])
+
         assert all_tokens == should_be_tokens
         assert len(all_tokens) <= max_length
 
@@ -106,8 +123,9 @@ def test_tokenize_dialogue_multi_turn_truncation_left(self, user_response_pairs,
         dialog = tokenize_dialogue(flat_convo, self.tokenizer, max_length=max_length)
 
         all_tokens = sum((dm.tokens for dm in dialog), ())
-
         should_be_tokens = sum(tokenized_flat_convo, ())[-max_length:]
-        assert all_tokens == should_be_tokens
+        if dialog[0] == DialogMessage(is_output=False, tokens=(self.tokenizer.eos_token_id,)):
+            should_be_tokens = (self.tokenizer.eos_token_id, *should_be_tokens[-max_length + 1 :])
 
+        assert all_tokens == should_be_tokens
         assert len(all_tokens) <= max_length

trlx/pipeline/offline_pipeline.py

@@ -26,15 +26,15 @@ class DialogMessage:
 
 
 def tokenize_dialogue(  # noqa: C901
-    dialogue: Union[str, List[str]], tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], max_length=2048
+    dialogue: Union[str, Iterable[str]], tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], max_length=2048
 ) -> List[DialogMessage]:
     """
     Tokenize sample with the interleaved form of (prompt_1, output_1, prompt_2, output_2...)
     """
     if isinstance(dialogue, str):
         bos_token = tokenizer.bos_token or tokenizer.eos_token
         dialogue = [bos_token, dialogue]
-    elif isinstance(dialogue, tuple):
+    elif isinstance(dialogue, Iterable):
         if len(dialogue) % 2 != 0:
             raise ValueError("Dialogue must have an even number of phrases, alternating prompt and output")
         dialogue = list(dialogue)
@@ -64,9 +64,17 @@ def tokenize_dialogue(  # noqa: C901
         truncated = [DialogMessage(is_output=m.is_output, tokens=m.tokens[::-1]) for m in truncated[::-1]]
 
     # remove empty messages
-    truncated = [t for t in truncated if len(t.tokens) > 0]
+    out = [t for t in truncated if len(t.tokens) > 0]
 
-    return truncated
+    if out[0].is_output:
+        if sum(map(lambda msg: len(msg.tokens), out)) == max_length:
+            if tokenizer.truncation_side == "left":
+                out[0].tokens = out[0].tokens[1:]
+            else:
+                out[-1].tokens = out[-1].tokens[:-1]
+
+        out.insert(0, DialogMessage(False, (tokenizer.bos_token_id,)))
+    return out
 
 
 class DialogStore(BaseRolloutStore):