An in-memory dataset for supervised finetuning with `input_ids` and `labels`. Args: data: A list of samples (dicts). The target/label must be stored under the key 'output' and the instruction or other data can be stored under any key as long as it is compatible with the give
| 6 | from lit_gpt.datamodules.typings.base import SFTDatasetItem |
| 7 | |
| 8 | class SFTDataset(Dataset): |
| 9 | """An in-memory dataset for supervised finetuning with `input_ids` and `labels`. |
| 10 | |
| 11 | Args: |
| 12 | data: A list of samples (dicts). The target/label must be stored under the key 'output' and the instruction |
| 13 | or other data can be stored under any key as long as it is compatible with the given prompt template. |
| 14 | tokenizer: The tokenizer to use. Should match the one that was used to pretrain the model. |
| 15 | prompt_style: The style to apply to prompts. See `litgpt.prompts` for a list of available styles. |
| 16 | max_seq_length: Truncate sequences that are longer than this value. By default, no truncation is applied. |
| 17 | mask_prompt: Whether to mask the prompt section from the label (with ``ignore_index``). |
| 18 | ignore_index: The index to use for elements to be ignored in the label. |
| 19 | transform: An optional transform to apply to the sample before it gets tokenized. Use this to rename the |
| 20 | keys in the dataset to the expected 'instruction' and 'output' keys. |
| 21 | |
| 22 | Returns a dict with two keys: |
| 23 | input_ids: The encoded prompt + response |
| 24 | labels: Same as input_ids, unless ``mask_prompt=True`` in which case the 'prompt' part is replaced with |
| 25 | the ``ignore_index``. |
| 26 | """ |
| 27 | |
| 28 | def __init__( |
| 29 | self, |
| 30 | data: List[Dict[str, str]], |
| 31 | tokenizer: Tokenizer, |
| 32 | prompt_style: Union[str, PromptStyle], |
| 33 | max_seq_length: int = -1, |
| 34 | mask_prompt: bool = True, |
| 35 | ignore_index: int = -100, |
| 36 | transform: Optional[Callable[[Any], Any]] = None, |
| 37 | ) -> None: |
| 38 | self.data = data |
| 39 | self.tokenizer = tokenizer |
| 40 | self.prompt_style = ( |
| 41 | prompt_style |
| 42 | if isinstance(prompt_style, PromptStyle) |
| 43 | else PromptStyle.from_name(prompt_style) |
| 44 | ) |
| 45 | self.max_seq_length = max_seq_length |
| 46 | self.mask_prompt = mask_prompt |
| 47 | self.ignore_index = ignore_index |
| 48 | self.transform = transform |
| 49 | |
| 50 | def __len__(self) -> int: |
| 51 | return len(self.data) |
| 52 | |
| 53 | def __getitem__(self, idx: int) -> SFTDatasetItem: |
| 54 | example = self.data[idx] |
| 55 | if self.transform is not None: |
| 56 | example = self.transform(example) |
| 57 | prompt = self.prompt_style.apply(prompt=example['instruction'], **example) |
| 58 | prompt_and_response = prompt + example['output'] |
| 59 | encoded_prompt = self.tokenizer.encode( |
| 60 | prompt, bos=True, max_length=self.max_seq_length |
| 61 | ) |
| 62 | encoded_prompt_and_response = self.tokenizer.encode( |
| 63 | prompt_and_response, bos=True, eos=True, max_length=self.max_seq_length |
| 64 | ) |
| 65 |
no outgoing calls