[feat] Support tying metadata to each prompt #421
Conversation
| @@ -282,7 +282,7 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0): # noq | |||
| exp_generate_time = time() | |||
|
|
|||
| # Generate samples from the language model (similar to using HuggingFace `generate` method) | |||
| samples = self.generate(**batch) | |||
| samples = self.generate(batch["input_ids"], batch["attention_mask"]) | |||
There was a problem hiding this comment.
Aside: This change makes it clear that next(self.prompt_iterator) doesn't materialize a PromptBatch instance; in fact, PromptBatch objects are never created anywhere. Sigh.
There was a problem hiding this comment.
Why don't we depreciate prompt batch then.
| original_samples = [p + o + reward_tokenizer.eos_token for p, o in zip(prompts, original_output)] | ||
| original_rewards = get_reward(original_samples) |
| @@ -125,7 +134,15 @@ def __len__(self) -> int: | |||
| return len(self.prompts) | |||
|
|
|||
| def create_loader(self, batch_size: int, shuffle=False) -> DataLoader: | |||
| collate_fn = DataCollatorWithPadding(self.tokenizer) if self.tokenizer else torch.vstack | |||
| def collate_fn(xs): | |||
| out = self.tokenizer.pad([{"input_ids": x["input_ids"]} for x in xs], return_tensors="pt") | |||
There was a problem hiding this comment.
Is it safe to assume self.tokenizer will always be available? Not sure if the previous check if self.tokenizer else torch.vstack is just stale code?
There was a problem hiding this comment.
It was a stale code and the single example which didn't use tokenizer at the time (ilql_randomwalks) was also converted to use tokenizer to drop these checks, out of which it is the last, for consistency; right now tokenizer is required in almost every other piece of code, so I'm not sure if that feature is missed.
| super().__init__() | ||
|
|
||
| if isinstance(prompts[0], dict): | ||
| metadata = prompts | ||
| prompts = [x.pop("prompt") for x in metadata] |
There was a problem hiding this comment.
In PromptPipeline's constructor docstring, we should document that "prompt" is an expected field of metadata when prompts consists of dictionaries.
This PR lets users pass an optional dictionary alongside each prompt into
PromptPipeline, which would get passed later, together with generated completions, into a reward function in a form of additional keywords per each key in the original dictionary.Fix of #297, #301, #340
https://wandb.ai/sorry/trlx-references/reports/add-prompts-metadata-v-main--Vmlldzo0MDAyMDIy