Entraîner un modèle linguistique causal à partir de zéro
Source :vignettes/examples/text-generation.Rmd
text-generation.RmdCet exemple est une adaptation du cours ‘Entraîner un modèle linguistique causal à partir de zéro’ du Hugging Face NLP course.
library(torch)
library(tok)
library(luz)
library(minhub) # remotes::install_github("mlverse/minhub")
#library(tidyverse)
options(arrow.skip_nul = TRUE)
library(arrow)Données
La première étape est d’implémenter un jeu de données torch qui rassemble des données et les pré-traite pour qu’elles soient au format approprié pour entraîner le modèle.
Cela signifie que nous devons :
- Télécharger les données textuelles
- Entraîner un tokeniseur pour ce jeu de données
- Être en mesure de produire des séquences de tokens dans le format attendu par le modèle
Nous allons utiliser 2 jeux de données disponibles sur le Hub d’Hugging Face. Le premier contient tous les codes sources des paquets R disponibles sur CRAN. Le second contient tous les codes R qui sont disponibles dans les dumps GitHub. Les deux jeux de données sont au format Parquet. Nous allons implémenter une fonction qui télécharge et cache les données, puis renvoie une seule table au format arrow contenant toutes les données.
read_dataset <- function(source) {
d <- source |>
hfhub::hub_snapshot(repo_type = "dataset", allow_patterns = "parquet$") |>
fs::path("data/r") |>
arrow::open_dataset() |>
dplyr::filter(stringr::str_detect(path, ".*\\.[rR]$")) |>
dplyr::select(content) |>
dplyr::mutate(content = arrow::cast(content, arrow::string())) |>
dplyr::filter(!is.na(content)) |>
dplyr::collect() %>%
# le jeu de données contient des caractères utf8 invalides...
# nous devons les supprimer, sinon on obtient une erreur des tokenizers
dplyr::filter(utf8::utf8_valid(content))
}
read_datasets <- function() {
dplyr::bind_rows(
read_dataset("dfalbel/cran-packages"),
read_dataset("dfalbel/github-r-repos")
)
}Ensuite, nous implémentons une fonction qui entraîne un tokeniseur sur notre jeu de données.
create_tokenizer <- function(text, vocab_size, special_tokens) {
tok <- tok::tokenizer$new(tok::model_bpe$new())
tok$pre_tokenizer <- tok::pre_tokenizer_byte_level$new(add_prefix_space = FALSE)
tok$decoder <- tok::decoder_byte_level$new()
tok$post_processor <- tok::processor_byte_level$new(trim_offsets = FALSE)
tok$train_from_memory(
text,
tok::trainer_bpe$new(vocab_size = vocab_size, special_tokens = special_tokens)
)
tok
}
# code de débogage pour le tokeniseur
# data <- read_datasets()
# tok <- create_tokenizer(data$content)Nous pouvons enfin implémenter le jeu de données torch que nous
allons utiliser pour entraîner le modèle. Nous allons utiliser
torch::iterable_dataset au lieu de
torch::dataset. La principale motivation est que nous ne
pouvons pas
vraiment savoir le nombre total d’échantillons dans le jeu de données,
donc nous pouvons implémenter une méthode .getitem() pour
obtenir n’importe quel échantillon arbitraire. Ainsi, nous implémentons
la méthode .iter qui retourne
un nouvel échantillon à chaque appel.
r_sources_dataset <- torch::iterable_dataset(
"r_sources_dataset",
initialize = function(root = ".", vocab_size = 20000, context_length = 128) {
self$data <- read_datasets()
self$context_length <- context_length
self$index <- sample.int(nrow(self$data))
# nous créons un tokeniseur que si celui-ci n'existe pas déjà, sinon nous le chargeons simplement
tok_path <- file.path(root, glue::glue("tokenizer-{vocab_size}.json"))
if (!file.exists(tok_path)) {
self$tok <- create_tokenizer(
as.character(self$data$content),
vocab_size,
c("<fbegin>", "<fend>")
)
fs::dir_create(root)
self$tok$save(tok_path)
} else {
self$tok <- tok::tokenizer$from_file(tok_path)
}
},
.iter = function() {
i <- 1L
sequence <- c()
function() {
while (length(sequence) < (self$context_length + 1) && i <= nrow(self$data)) {
sequence <<- c(
sequence,
self$tok$encode(paste("<fbegin>", as.character(self$data$content[self$index[i]]), "<fend>"))$ids
)
i <- i + 1L
}
if (length(sequence) < (self$context_length + 1)) {
return(coro::exhausted())
}
on.exit({
sequence <<- sequence[-seq_len(self$context_length)]
})
list(
input_ids = sequence[seq_len(self$context_length)] + 1L,
labels = sequence[2:(self$context_length + 1)] + 1L
)
}
}
)
# code de débogage pour le jeu de données
# ds <- r_sources_dataset("~/Downloads/")
# it <- ds$.iter()
# it()
# ds$tok$get_vocab_size()Ce jeu de données est bien trop volumineux pour entraîner le modèle sur tous les documents dans cet exemple. Il est également difficile de prédire combien de temps cela prendra d’arriver à la fin de l’entraînement. Pour simplifier, nous définissons un jeu de données itérable utilisé pour exécuter le jeu de données ci-dessus pendant un nombre fixe d’étapes. Ce n’est pas nécessaire, mais rend l’utilisation de luz plus agréable, car nous pouvons facilement définir combien de tokens nous voulons entraîner notre modèle.
fixed_steps_iterable_dataset <- iterable_dataset(
"fixed_steps_dataset",
initialize = function(dataset, steps) {
self$dataset <- dataset
self$steps <- steps
},
.iter = function() {
i <- 1L
iter <- NULL
function() {
if (i > self$steps) {
return(coro::exhausted())
}
i <<- i + 1L
if (is.null(iter) || coro::is_exhausted(data <- iter())) {
iter <<- self$dataset$.iter()
data <- iter()
}
data
}
},
.length = function() {
self$steps
}
)Maintenant, nous pouvons définir le modèle que nous allons entraîner.
Nous utiliserons une version légère de GPT2. Nous définissons également
une méthode generate nous permettant d’échantillonner à
partir du modèle donné un contexte initial.
net <- nn_module(
initialize = function() {
self$gpt <- minhub::gpt2(
vocab_size = 20000,
pdrop = 0.1
)
},
forward = function(x) {
self$gpt(x)$transpose(2,3)
},
generate = function(x, temperature = 1, iter = 50, top_k = 10) {
# échantillonne à partir du modèle given un vecteur de contexte.
for (i in seq_len(iter)) {
logits <- self$forward(x)[,,-1]
logits <- logits/temperature
c(prob, ind) %<-% logits$topk(top_k)
logits <- torch_full_like(logits, -Inf)$scatter_(-1, ind, prob)
logits <- nnf_softmax(logits, dim = -1)
id_next <- torch_multinomial(logits, num_samples = 1)
x <- torch_cat(list(x, id_next), dim = 2)
}
x
}
)
# code de débogage pour le modèle
# ds <- torch::dataloader(r_sources_dataset("~/Downloads/"), batch_size = 32)
# batch <- coro::collect(ds, 1)[[1]]
# str(batch)
# m <- net()
# str(m(batch$input_ids))Pour faciliter l’inspection de la formation, nous définirons également un callback qui imprime un échantillon du modèle à chaque époque.
# échantillonne à partir du modèle en utilisant le contexte.
generate <- function(model, tok, context, ...) {
local_no_grad() # désactive les gradients pour l'échantillonnage
x <- tok$encode(context)$ids + 1L
x <- torch_tensor(x)[NULL,]$to(device = model$device)
content <- as.integer(model$generate(x, ...)$cpu())
tok$decode(content - 1L)
}
display_cb <- luz_callback(
initialize = function() {},
on_epoch_end = function() {
local_no_grad()
# sample from the model...
context <- "# creates a linear model"
text <- generate(ctx$model, dataset$dataset$tok, context, iter = 100)
cli::cli_rule()
cat(text, "\n")
cli::cli_rule()
}
)Nous pouvons maintenant entraîner le modèle. Nous définissons un entraînement du modèle pour un demi-milliard de tokens pendant un total de 100 époques.
n_tokens <- 500e6
batch_size <- 16
epochs <- 100
context_length <- 256L
steps <- n_tokens / context_length / epochs
dataset <- fixed_steps_iterable_dataset(
r_sources_dataset(context_length = context_length),
steps = steps
)
fitted <- net %>%
setup(
optimizer = optim_adam,
loss = nn_cross_entropy_loss()
) %>%
set_opt_hparams(lr = 3e-4) |>
fit(
dataset,
epochs = epochs,
dataloader_options = list(batch_size = batch_size),
callbacks = list(
luz_callback_lr_scheduler(
torch::lr_one_cycle,
max_lr = 0.1,
epochs = epochs,
steps_per_epoch = steps/batch_size,
call_on = "on_batch_end"
),
luz_callback_gradient_clip(max_norm = 1),
display_cb()
),
verbose = TRUE
)
luz::luz_save(fitted, "model.pt")Ensuite, nous pouvons utiliser le modèle pour générer du texte en fonction d’un prompt avec :