nemo_rl.data.datasets

Contenu du Module

Classes

Fonctions

Données

TokenizerType

API

nemo_rl.data.datasets.TokenizerType

Valeur : None

class nemo_rl.data.datasets.AllTaskProcessedDataset(dataset: datasets.Dataset | typing.Any, tokenizer: nemo_rl.data.datasets.TokenizerType, default_task_data_spec: nemo_rl.data.interfaces.TaskDataSpec, task_data_processors: dict[str, tuple[nemo_rl.data.interfaces.TaskDataSpec, nemo_rl.data.interfaces.TaskDataProcessFnCallable]] | nemo_rl.data.interfaces.TaskDataProcessFnCallable, max_seq_length: typing.Optional[int] = None)

Ensemble de données pour le traitement de données mono-tâche ou multi-tâches avec tokenisation et traitement spécifiques à la tâche.

Args : dataset : Ensemble de données d’entrée contenant des données brutes tokenizer : Tokenizer pour le traitement du texte default_task_data_spec : Spécifications de traitement de tâche par défaut. Dans le cas d’une tâche unique, il s’agit de la spécification utilisée pour traiter toutes les entrées. Dans le cas de tâches multiples, toutes les valeurs non spécifiées dans les spécifications spécifiques à la tâche seront extraites de la spécification par défaut. task_data_processors : Soit un seul TaskDataProcessFnCallable pour une tâche unique, soit un dictionnaire associant des noms de tâches à (TaskDataSpec, TaskDataProcessFnCallable) pour des tâches multiples max_seq_length : Longueur maximale de séquence pour les sorties tokenisées

__len__() -> int

encode_single(text: typing.Union[str, list[str]]) -> tuple[list[int] | torch.Tensor, int]

Prend soit une chaîne unique, soit une liste de chaînes représentant plusieurs tours de la même conversation.

Retourne une liste unique (concaténée) d’ID tokenisés et la longueur des ID tokenisés.

__getitem__(idx: int) -> nemo_rl.data.interfaces.DatumSpec

Retourne une invite unique.

nemo_rl.data.datasets.rl_collate_fn(data_batch: list[nemo_rl.data.interfaces.DatumSpec]) -> nemo_rl.distributed.batched_data_dict.BatchedDataDict[typing.Any]

Fonction de collation pour l’entraînement par renforcement.

nemo_rl.data.datasets.eval_collate_fn(data_batch: list[nemo_rl.data.interfaces.DatumSpec]) -> nemo_rl.distributed.batched_data_dict.BatchedDataDict[typing.Any]

Fonction de collation pour l’évaluation.

Prend une liste d’échantillons de données et les combine en un seul dictionnaire de lots pour l’évaluation du modèle.

Args : data_batch : Liste d’échantillons de données avec les champs message_log, extra_env_info et idx.

Retourne : BatchedDataDict avec les champs message_log, extra_env_info et idx.

Exemples :

>>> import torch
>>> from nemo_rl.data.datasets import eval_collate_fn
>>> from nemo_rl.data.interfaces import DatumSpec
>>> data_batch = [
...     DatumSpec(
...         message_log=[{"role": "user", "content": "Hello", "token_ids": torch.tensor([1, 2, 3])}],
...         extra_env_info={'ground_truth': '1'},
...         idx=0,
...     ),
...     DatumSpec(
...         message_log=[{"role": "assistant", "content": "Hi there", "token_ids": torch.tensor([4, 5, 6, 7])}],
...         extra_env_info={'ground_truth': '2'},
...         idx=1,
...     ),
... ]
>>> output = eval_collate_fn(data_batch)
>>> output['message_log'][0]
[{'role': 'user', 'content': 'Hello', 'token_ids': tensor([1, 2, 3])}]
>>> output['message_log'][1]
[{'role': 'assistant', 'content': 'Hi there', 'token_ids': tensor([4, 5, 6, 7])}]
>>> output['extra_env_info']
[{'ground_truth': '1'}, {'ground_truth': '2'}]
>>> output['idx']
[0, 1]

nemo_rl.data.datasets.preference_collate_fn(
    data_batch: list[nemo_rl.data.interfaces.DPODatumSpec],
    tokenizer: nemo_rl.data.datasets.TokenizerType,
    make_sequence_length_divisible_by: int,
    add_loss_mask: bool
) -> nemo_rl.distributed.batched_data_dict.BatchedDataDict[typing.Any]

Fonction de collation pour l’entraînement sur des données de préférence.

Cette fonction sépare les réponses choisies et rejetées pour créer deux exemples par invite. Les exemples choisis et rejetés sont entrelacés le long de la dimension du lot, ce qui donne un lot de taille 2 * len(data_batch).

Args : data_batch : Liste d’échantillons de données avec les champs message_log_chosen, message_log_rejected, length_chosen, length_rejected, loss_multiplier, idx et task_name. tokenizer : Tokenizer pour le traitement du texte make_sequence_length_divisible_by : Rendre la longueur de séquence divisible par cette valeur add_loss_mask : Indique s’il faut ajouter un token_mask aux données renvoyées Retourne : BatchedDataDict avec les champs input_ids, input_lengths, token_mask (optionnel) et sample_mask.

nemo_rl.data.datasets.assert_no_double_bos(
    token_ids: torch.Tensor, tokenizer: nemo_rl.data.datasets.TokenizerType
) -> None

Vérifie qu’il n’y a pas de double jeton de début (BOS) dans le message.

Args : token_ids : Liste d’ID de jetons tokenizer : Tokenizer