API de décodage

Voir sur TensorFlow.org Exécuter dans Google Colab Voir sur GitHub Télécharger le cahier

Aperçu

Dans un passé récent, il y a eu beaucoup de recherches sur la génération de langage avec des modèles auto-régressifs. Dans la génération de langage auto-régressif, la distribution de probabilité de jeton au moment de l' étape K dépend de jeton-prédictions du modèle jusqu'à l' étape K-1. Pour ces modèles, les stratégies de décodage telles que la recherche de faisceau, Greedy, Top-p et Top-k sont des composantes essentielles du modèle et influencent largement le style / nature de la sortie générée jeton à un K pas de temps donné.

Par exemple, la recherche de faisceau réduit le risque de passer à côté des jetons de haute probabilité cachés en gardant le plus num_beams probables des hypothèses à chaque pas de temps et finalement choisir l'hypothèse qui a la plus forte probabilité globale. Murray et al. (2018) et Yang et al. (2018) montrent que la recherche de faisceau fonctionne bien dans les tâches de traduction automatique. Les deux stratégies de recherche Beam & Greedy ont une possibilité de générer des jetons de répétition.

Fan et. al (2018) a présenté l' échantillonnage Top-K, où K jetons les plus probables sont filtrés et la masse de probabilité est redistribuée parmi seuls les jetons K.

Ari Holtzman et. al (2019) a présenté l' échantillonnage Haut-p, qui choisit le plus petit possible l' ensemble de jetons avec une probabilité cumulative qui ajoute jusqu'à la probabilité p. La masse de probabilité est alors redistribuée parmi cet ensemble. De cette façon, la taille de l'ensemble de jetons peut augmenter et diminuer de manière dynamique. Top-p, Top-k sont généralement utilisés dans des tâches telles que l' histoire génération.

L'API de décodage fournit une interface pour expérimenter différentes stratégies de décodage sur des modèles auto-régressifs.

  1. Les stratégies d'échantillonnage suivantes sont fournies dans sampling_module.py, qui hérite de la classe Decoding de base :

  2. La recherche de faisceau est fournie dans beam_search.py. github

Installer

pip install -q -U tensorflow-text
pip install -q tf-models-nightly
import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf

from official import nlp
from official.nlp.modeling.ops import sampling_module
from official.nlp.modeling.ops import beam_search
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/pkg_resources/__init__.py:119: PkgResourcesDeprecationWarning: 0.18ubuntu0.18.04.1 is an invalid version and will not be supported in a future release
  PkgResourcesDeprecationWarning,

Initialiser le module d'échantillonnage dans TF-NLP.

  • symbols_to_logits_fn: Utilisez cette fermeture pour appeler le modèle pour prédire les logits pour l' index+1 étape. Les entrées et sorties pour cette fermeture sont les suivantes :
Args:
  1] ids : Current decoded sequences. int tensor with shape (batch_size, index + 1 or 1 if padded_decode is True)],
  2] index [scalar] : current decoded step,
  3] cache [nested dictionary of tensors] : Only used for faster decoding to store pre-computed attention hidden states for keys and values. More explanation in the cell below.
Returns:
  1] tensor for next-step logits [batch_size, vocab]
  2] the updated_cache [nested dictionary of tensors].

Le cache est utilisé pour un décodage plus rapide. Voici une référence mise en œuvre pour la fermeture ci - dessus.

  • length_normalization_fn: Utilisez cette fermeture pour revenir paramètre de normalisation de longueur.
Args: 
  1] length : scalar for decoded step index.
  2] dtype : data-type of output tensor
Returns:
  1] value of length normalization factor.
  • vocab_size: taille du vocabulaire de sortie.

  • max_decode_length: Scalar pour le nombre total d'étapes de décodage.

  • eos_id: Décodage arrêtera si tous les ids décodé de sortie du lot ont ce eos_id.

  • padded_decode: Réglez ce paramètre sur True si en cours d' exécution sur TPU. Les tenseurs sont remplis à max_decoding_length si cela est vrai.

  • top_k: top_k est activé si cette valeur est> 1.

  • top_p: top_p est activé si cette valeur est> 0 et <1.0

  • sampling_temperature: Il est utilisé pour réestimer la sortie softmax. La température fausse la distribution vers les jetons à haute probabilité et abaisse la masse dans la distribution de la queue. La valeur doit être positive. La basse température équivaut à gourmande et rend la distribution plus nette, tandis que la haute température la rend plus plate.

  • enable_greedy: Par défaut, cela est vrai et le décodage avide est activée. Pour expérimenter d'autres stratégies, veuillez définir ce paramètre sur False.

Initialiser les hyperparamètres du modèle

params = {}
params['num_heads'] = 2
params['num_layers'] = 2
params['batch_size'] = 2
params['n_dims'] = 256
params['max_decode_length'] = 4

Dans les architectures auto-régressif comme transformateur à base COdeurDECodeur modèles, le cache est utilisé pour le décodage séquentiel rapide. Il s'agit d'un dictionnaire imbriqué stockant des états cachés pré-calculés (clé et valeurs dans les blocs d'auto-attention et dans les blocs d'attention croisée) pour chaque couche.

Initialiser le cache.

cache = {
    'layer_%d' % layer: {
        'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], int(params['n_dims']/params['num_heads'])], dtype=tf.float32),
        'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], int(params['n_dims']/params['num_heads'])], dtype=tf.float32)
        } for layer in range(params['num_layers'])
    }
print("cache key shape for layer 1 :", cache['layer_1']['k'].shape)
cache key shape for layer 1 : (2, 4, 2, 128)

Définissez la fermeture pour la normalisation de la longueur si nécessaire.

Ceci est utilisé pour normaliser les scores finaux des séquences générées et est facultatif

def length_norm(length, dtype):
  """Return length normalization factor."""
  return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)

Créer model_fn

Dans la pratique, il sera remplacé par une mise en œuvre du modèle actuel, comme ici

Args:
i : Step that is being decoded.
Returns:
  logit probabilities of size [batch_size, 1, vocab_size]
probabilities = tf.constant([[[0.3, 0.4, 0.3], [0.3, 0.3, 0.4],
                              [0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                            [[0.2, 0.5, 0.3], [0.2, 0.7, 0.1],
                              [0.1, 0.1, 0.8], [0.1, 0.1, 0.8]]])
def model_fn(i):
  return probabilities[:, i, :]

Initialiser symboles_to_logits_fn

def _symbols_to_logits_fn():
  """Calculates logits of the next tokens."""
  def symbols_to_logits_fn(ids, i, temp_cache):
    del ids
    logits = tf.cast(tf.math.log(model_fn(i)), tf.float32)
    return logits, temp_cache
  return symbols_to_logits_fn

Cupide

Décodage Greedy sélectionne l'identifiant de jeton avec la probabilité la plus élevée en tant que son prochain id: \(id_t = argmax_{w}P(id | id_{1:t-1})\) à chaque pas de temps \(t\). Le croquis suivant montre un décodage gourmand.

greedy_obj = sampling_module.SamplingModule(
    length_normalization_fn=None,
    dtype=tf.float32,
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    vocab_size=3,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    padded_decode=False)
ids, _ = greedy_obj.generate(
    initial_ids=tf.constant([9, 1]), initial_cache=cache)
print("Greedy Decoded Ids:", ids)
Greedy Decoded Ids: tf.Tensor(
[[9 1 2 2 2]
 [1 1 1 2 2]], shape=(2, 5), dtype=int32)

top_k échantillonnage

Dans l' échantillonnage Top-K, K le plus probable sont filtrés et la masse de probabilité est redistribuée ids prochain jeton parmi seulement les K ids.

top_k_obj = sampling_module.SamplingModule(
    length_normalization_fn=length_norm,
    dtype=tf.float32,
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    vocab_size=3,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    sample_temperature=tf.constant(1.0),
    top_k=tf.constant(3),
    padded_decode=False,
    enable_greedy=False)
ids, _ = top_k_obj.generate(
    initial_ids=tf.constant([9, 1]), initial_cache=cache)
print("top-k sampled Ids:", ids)
top-k sampled Ids: tf.Tensor(
[[9 1 0 2 2]
 [1 0 1 2 2]], shape=(2, 5), dtype=int32)

top_p échantillonnage

Au lieu de l' échantillonnage seulement de la plus probable K jeton ids, dans le Top-p échantillonnage choisit de le plus petit ensemble possible de ids dont la probabilité cumulée est supérieure à la probabilité p.

top_p_obj = sampling_module.SamplingModule(
    length_normalization_fn=length_norm,
    dtype=tf.float32,
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    vocab_size=3,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    sample_temperature=tf.constant(1.0),
    top_p=tf.constant(0.9),
    padded_decode=False,
    enable_greedy=False)
ids, _ = top_p_obj.generate(
    initial_ids=tf.constant([9, 1]), initial_cache=cache)
print("top-p sampled Ids:", ids)
top-p sampled Ids: tf.Tensor(
[[9 1 1 2 2]
 [1 1 1 0 2]], shape=(2, 5), dtype=int32)

Décodage de recherche de faisceau

La recherche de faisceau réduit le risque de manquer des identifiants de jeton cachés à haute probabilité en conservant le nombre de faisceaux d'hypothèses les plus probables à chaque pas de temps et en choisissant éventuellement l'hypothèse qui a la probabilité globale la plus élevée.

beam_size = 2
params['batch_size'] = 1
beam_cache = {
    'layer_%d' % layer: {
        'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32),
        'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32)
        } for layer in range(params['num_layers'])
    }
print("cache key shape for layer 1 :", beam_cache['layer_1']['k'].shape)
ids, _ = beam_search.sequence_beam_search(
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    initial_ids=tf.constant([9], tf.int32),
    initial_cache=beam_cache,
    vocab_size=3,
    beam_size=beam_size,
    alpha=0.6,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    padded_decode=False,
    dtype=tf.float32)
print("Beam search ids:", ids)
cache key shape for layer 1 : (1, 4, 2, 256)
Beam search ids: tf.Tensor(
[[[9 0 1 2 2]
  [9 1 2 2 2]]], shape=(1, 2, 5), dtype=int32)