crscardellino's picture
Added device to the steps
f781de0
raw
history blame
3.89 kB
"""
Utilities file needed for the presentation.
flisol-cba-martin-fierro
Copyright (C) 2023 Cristian Cardellino
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
from datasets import DatasetDict
from transformers import PreTrainedTokenizerBase
from typing import Callable, Dict, List
def tokenize(
tokenizer: PreTrainedTokenizerBase, end_char: str = "\n"
) -> Callable[[Dict[str, List[str]]], DatasetDict]:
"""
Helper function that returns a function to use with the `map` method of
datasets.DatasetDict. It takes a tokenizer and generates a function that
applies that tokenizer with an optional `end_char` parameter (e.g. a
newline) that might be needed (e.g. when trying to tokenize and keep the
structure of a poem wich needs the newline after each sentence). This is
needed since the function `datasets.load_dataset` forcibly removes the
newlines characters.
Parameters
----------
tokenizer : PreTrainedTokenizerBase
The tokenizer to use for the tokenization process.
end_char : str
The end character to append to each line.
Returns
-------
Callable[[Dict[str, List[str]]], DatasetDict]
The function in charge of the tokenization process.
"""
def _tokenize(examples: Dict[str, List[str]]) -> DatasetDict:
return tokenizer([f"{e}{end_char}" for e in examples["text"]])
return _tokenize
def group_texts(examples: Dict[str, List[int]], block_size: int = 128) -> Dict[str, List[int]]:
"""
Helper function to concatenate a tokenized dataset (with the function above)
in chunks of `block_size`. The code was taken from
https://github.com/huggingface/notebooks/blob/main/examples/language_modeling.ipynb
Parameters
----------
examples : Dict[str, List[int]]
This is actually a `LazyBatch` from the transformers library, that is
given by the `DatasetDict.map` method. It should be the dataset returned
after tokenization with the function returned by `tokenize`. It should
have 2 main keys: 'input_ids' and 'attention_mask'.
block_size : int
The size of the block to use in the training process. If the total lenght
of the group of texts is not divisible by the block size it will ignore the
remaining data for simplicity.
Returns
-------
Dict[str, List[str, int]]
The dictionary that will provide the new dataset divided in chunks of
`block_size`.
"""
# Concatenate all texts.
concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
total_length = len(concatenated_examples["input_ids"])
# We drop the small remainder, we could add padding if the model supported
# it instead of this drop, you can customize this part to your needs
total_length = (total_length // block_size) * block_size
# Split by chunks of block_size length
result = {
k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
for k, t in concatenated_examples.items()
}
# labels to be used by the training phase, it copies since the Transformers
# library will be in charge of making the shift to the right
result["labels"] = result["input_ids"].copy()
return result