Feat/open instruct example (#381)

* added open-instruct

* fixed hooks

* GRPO

---------

Co-authored-by: Andreas Koepf <andreas.koepf@provisio.com>

examples/open-instruct/src/utils.py

@@ -0,0 +1,104 @@
+from typing import List
+
+import torch
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer
+
+import reasoning_gym
+from reasoning_gym.utils import SYSTEM_PROMPTS
+
+
+def list_preserving_collate(batch):
+    """
+    Custom collate function that preserves lists instead of converting to tensors.
+    """
+    token_ids = [item[0] for item in batch]
+    items = [item[1] for item in batch]
+    return token_ids, items
+
+
+class ReasoningGymDataset(Dataset):
+    def __init__(self, dataset_name, seed, size, tokenizer, developer_role, developer_prompt):
+        self.data = reasoning_gym.create_dataset(dataset_name, seed=seed, size=size)
+        self.tokenizer = tokenizer
+        self.developer_role = developer_role
+        self.developer_prompt = developer_prompt
+
+    def __len__(self):
+        return self.data.size
+
+    def __getitem__(self, index):
+        chat_message = [{"role": self.developer_role, "content": self.developer_prompt}]
+        item = self.data[index]
+        chat_message.append({"role": "user", "content": item["question"]})
+        prompt_text = self.tokenizer.apply_chat_template(chat_message, tokenize=True, add_generation_prompt=True)
+        prompt_text = [token for token in prompt_text if token != self.tokenizer.pad_token_id]
+        return prompt_text, item
+
+
+def pack_sequences(
+    queries: List[List[int]],
+    responses: List[List[int]],
+    pack_length: int,
+    pad_token_id: int,
+) -> "PackedSequences":
+    # assert padding token does not exist in queries and responses
+    query_responses = []
+    attention_masks = []
+    response_masks = []
+    num_actions = []
+    packed_seq_lens = []
+    cur_data = []
+    cur_response_mask = []
+    cur_num_actions = []
+    cur_packed_seq_lens = []
+    cur_attention_mask = []
+    offset = 0
+
+    for i in range(len(queries)):
+        query = queries[i]
+        response = responses[i]
+        # remove padding (but using vllm so this should not be needed, but just in case)
+        query = [t for t in query if t != pad_token_id]
+        response = [t for t in response if t != pad_token_id]
+        query_response = query + response
+
+        if len(query_response) + len(cur_data) > pack_length:
+            query_responses.append(cur_data)
+            response_masks.append(cur_response_mask)
+            attention_masks.append(cur_attention_mask)
+            num_actions.append(cur_num_actions)
+            packed_seq_lens.append(cur_packed_seq_lens)
+            cur_data = []
+            cur_response_mask = []
+            cur_attention_mask = []
+            cur_num_actions = []
+            cur_packed_seq_lens = []
+            offset = i
+
+        cur_data.extend(query_response)
+        cur_num_actions.append(len(response))
+        cur_packed_seq_lens.append(len(query_response))
+        cur_response_mask.extend([0 for _ in range(len(query))] + [i + 1 for _ in range(len(response))])
+        cur_attention_mask.extend([i + 1 - offset for _ in range(len(query_response))])
+
+    if len(cur_data) > 0:
+        query_responses.append(cur_data)
+        response_masks.append(cur_response_mask)
+        attention_masks.append(cur_attention_mask)
+        num_actions.append(cur_num_actions)
+        packed_seq_lens.append(cur_packed_seq_lens)
+
+    attention_masks_list = [torch.tensor(t) for t in attention_masks]
+
+    from open_instruct.rl_utils2 import PackedSequences, reset_position_ids
+
+    return PackedSequences(
+        query_responses=[torch.tensor(t) for t in query_responses],
+        attention_masks=attention_masks_list,
+        position_ids=[reset_position_ids(t.unsqueeze(0)).squeeze(0) for t in attention_masks_list],
+        response_masks=[torch.tensor(t) for t in response_masks],
+        original_responses=responses,
+        num_actions=[torch.tensor(t) for t in num_actions],
+        packed_seq_lens=[torch.tensor(t) for t in packed_seq_lens],
+    )