nemo_rl.data.hf_datasets.oai_format_dataset

Contenu du module

Classes

API

class nemo_rl.data.hf_datasets.oai_format_dataset.OpenAIFormatDataset(train_ds_path: str, val_ds_path: str, chat_key: str = 'messages', system_key: str | None = None, system_prompt: str | None = None)

Cette classe est utilisée pour charger un jeu de données SFT au format OpenAI.

Le jeu de données doit être au format suivant : { “messages”: [ {“role”: “system”, “content”: “Vous êtes un assistant utile.”}, {“role”: “user”, “content”: “Quelle est la capitale de la France ?”}, {“role”: “assistant”, “content”: “La capitale de la France est Paris.”} ] } system_key et system_prompt sont facultatifs. S’ils sont fournis, ils seront ajoutés au début du jeu de données. chat_key doit être la clé de la liste des messages. Les conversations à plusieurs tours sont prises en charge. Le dernier message de la conversation doit provenir de l’assistant.

add_messages_key(example: dict[str, typing.Any]) -> dict[str, list[dict[str, typing.Any]]]