API di decodifica

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza su GitHub Scarica taccuino

Panoramica

Nel recente passato, c'è stata molta ricerca sulla generazione del linguaggio con modelli auto-regressivi. Nella generazione lingua auto-regressivo, la distribuzione di probabilità di gettone al passo temporale K dipende token-le previsioni del modello fino passo K-1. Per questi modelli, strategie decodifica come la ricerca del fascio, avido, Top-p, e Top-k sono componenti critici del modello e in gran parte influenzano lo stile / natura del output generato il token in un momento determinato passo K.

Ad esempio, la ricerca del fascio riduce il rischio di perdere nascosti gettoni alta probabilità di mantenere il maggior numero di num_beams probabili ipotesi ad ogni passo e alla fine la scelta l'ipotesi che ha la più alta probabilità complessiva. Murray et al. (2018) e Yang et al. (2018) mostrano che la ricerca fascio funziona bene in compiti di traduzione automatica. Sia fascio di ricerca e strategie Greedy hanno una possibilità di generare token ripetizione.

Fan et. al (2018) ha introdotto Top-K di campionamento, in cui il K più probabili gettoni sono filtrati e la massa di probabilità è solo ridistribuiti tra quelle pedine K.

Ari Holtzman et. al (2019) ha introdotto Top-p di campionamento, che sceglie dal più piccolo possibile di gettoni con probabilità cumulativa che aggiunge fino la probabilità p. La massa di probabilità viene quindi ridistribuita tra questo insieme. In questo modo, la dimensione del set di token può aumentare e diminuire dinamicamente. Top-p, Top-k sono generalmente utilizzati in compiti come story-generazione.

L'API di decodifica fornisce un'interfaccia per sperimentare diverse strategie di decodifica su modelli auto-regressivi.

  1. Le seguenti strategie di campionamento sono fornite in sampling_module.py, che eredita dalla classe Decoding di base:

  2. La ricerca del raggio è fornita in beam_search.py. github

Impostare

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,

Inizializza il modulo di campionamento in TF-NLP.

  • symbols_to_logits_fn: Utilizzare questa chiusura per richiamare il modello per prevedere i logit per l' index+1 passo. Gli ingressi e le uscite per questa chiusura sono i seguenti:
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].

La cache viene utilizzata per una decodifica più rapida. Qui è un riferimento di attuazione per la chiusura sopra.

  • length_normalization_fn: Utilizzare questa chiusura per il ritorno parametro di lunghezza normalizzazione.
Args: 
  1] length : scalar for decoded step index.
  2] dtype : data-type of output tensor
Returns:
  1] value of length normalization factor.

  • vocab_size: vocabolario di uscita.

  • max_decode_length: scalare per numero totale di passi decodifica.

  • eos_id: decodifica si arresta se tutti gli ID di uscita decodificato in batch hanno questo eos_id.

  • padded_decode: Impostare su True se in esecuzione su TPU. I tensori vengono riempiti fino a max_decoding_length se questo è True.

  • top_k: top_k è abilitata se questo valore è> 1.

  • top_p: top_p è attivato se questo valore è> 0 e <1,0

  • sampling_temperature: Questo è usato per ri-stimare la produzione SoftMax. La temperatura distorce la distribuzione verso token ad alta probabilità e riduce la massa nella distribuzione della coda. Il valore deve essere positivo. La bassa temperatura equivale a greedy e rende la distribuzione più nitida, mentre l'alta temperatura la rende più piatta.

  • enable_greedy: per impostazione predefinita, questo è vero e la decodifica avido è abilitato. Per sperimentare altre strategie, impostalo su False.

Inizializzare gli iperparametri del modello

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

Nelle architetture auto-regressivo come base di Transformer Encoder-Decoder modelli, cache viene utilizzato per la decodifica sequenziale veloce. È un dizionario nidificato che memorizza stati nascosti precalcolati (chiave e valori nei blocchi di attenzione personale e nei blocchi di attenzione incrociata) per ogni livello.

Inizializza la 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)

Se necessario, definire la chiusura per la normalizzazione della lunghezza.

Questo è usato per normalizzare i punteggi finali delle sequenze generate ed è facoltativo

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

Crea modello_fn

In pratica, questo sarà sostituito da un'implementazione effettiva modello come qui

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, :]

Inizializza simboli_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

Avido

Decodifica avido seleziona l'id token con la più alta probabilità come sua prossima id: \(id_t = argmax_{w}P(id | id_{1:t-1})\) ad ogni passo temporale \(t\). Lo schizzo seguente mostra la decodifica greedy.

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)

campionamento top_k

Nel campionamento Top-K, la K molto probabilmente gli ID token successivo vengono filtrati e la massa di probabilità viene ridistribuita tra solo gli ID K.

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 campionamento

Invece di campionamento solo dal più probabile K di token ID, in Top-p campionamento sceglie tra il più piccolo possibile di ID la cui probabilità cumulativa supera 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)

Decodifica della ricerca del raggio

La ricerca fascio riduce il rischio di perdere gli ID token nascosti ad alta probabilità mantenendo i num_beam di ipotesi più probabili in ogni fase temporale ed eventualmente scegliendo l'ipotesi che ha la probabilità complessiva più alta.

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)