API de decodificação

Ver no TensorFlow.org Executar no Google Colab Ver no GitHub Baixar caderno

Visão geral

No passado recente, houve muitas pesquisas na geração de linguagem com modelos auto-regressivos. Na geração de linguagem auto-regressivo, a distribuição de probabilidade de sinal no instante K é dependente de token-previsões do modelo até o passo K-1. Para estes modelos, estratégias de decodificação, como a busca Beam, Greedy, Top-p, e Top-k são componentes críticos do modelo e influenciar em grande parte o estilo / natureza da saída gerada simbólico de um tempo determinado passo K.

Por exemplo, a pesquisa Feixe reduz o risco de falta escondidos fichas alta probabilidade de manter a maioria dos num_beams prováveis hipóteses em cada passo de tempo e, eventualmente, escolher a hipótese de que tem a maior probabilidade geral. Murray et al. (2018) e Yang et al. (2018) mostram que busca feixe funciona bem em tarefas de tradução automática. Ambos pesquisa Feixe e estratégias Greedy ter uma possibilidade de gerar símbolos de repetição.

Fan et. al (2018) introduziu amostragem Top-K, em que K o fichas mais prováveis são filtrados e a massa de probabilidade é redistribuído entre somente os símbolos K.

Ari Holtzman et. al (2019) introduziu Top-p amostragem, que escolhe desde o mais pequeno possível conjunto de fichas com probabilidade cumulativa que adiciona até a probabilidade p. A massa de probabilidade é então redistribuída entre este conjunto. Dessa forma, o tamanho do conjunto de tokens pode aumentar ou diminuir dinamicamente. Top-p, Top-k são geralmente usados em tarefas como a história geração.

A API de decodificação fornece uma interface para experimentar diferentes estratégias de decodificação em modelos auto-regressivos.

  1. As seguintes estratégias de amostragem são fornecidas em sampling_module.py, que herda da classe Decoding base:

  2. A pesquisa de feixe é fornecida em beam_search.py. github

Configurar

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,

Inicialize o Módulo de Amostragem no TF-NLP.

  • symbols_to_logits_fn: Use este encerramento para chamar o modelo para prever as logits para o index+1 etapa. As entradas e saídas para este fechamento são as seguintes:
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].

O cache é usado para uma decodificação mais rápida. Aqui é uma referência de implementação para o fechamento acima.

  • length_normalization_fn: Use este encerramento para o regresso de parâmetro comprimento de normalização.
Args: 
  1] length : scalar for decoded step index.
  2] dtype : data-type of output tensor
Returns:
  1] value of length normalization factor.
  • vocab_size: tamanho do vocabulário de saída.

  • max_decode_length: escalar para o número total de passos de descodificação.

  • eos_id: Decoding vai parar se todos os ids de saída decodificados no lote tem esse eos_id.

  • padded_decode: Defina como True se em execução em TPU. Os tensores são preenchidos com max_decoding_length se for True.

  • top_k: top_k está habilitado se este valor for> 1.

  • top_p: top_p está habilitado se este valor for> 0 e <1,0

  • sampling_temperature: Este é utilizado para re-estimar a saída Softmax. A temperatura inclina a distribuição em direção a tokens de alta probabilidade e reduz a massa na distribuição da cauda. O valor tem que ser positivo. A baixa temperatura equivale a guloso e torna a distribuição mais nítida, enquanto a alta temperatura a torna mais plana.

  • enable_greedy: Por padrão, isso é verdade e decodificação ganancioso está habilitado. Para experimentar outras estratégias, defina como False.

Inicialize os hiperparâmetros do modelo

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

Em arquiteturas auto-regressivos como base Transformer Encoder-Decoder modelos, Cache é usado para rápida decodificação seqüencial. É um dicionário aninhado que armazena estados ocultos pré-calculados (chave e valores nos blocos de autoatenção e nos blocos de atenção cruzada) para cada camada.

Inicialize o 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)

Defina o fechamento para normalização do comprimento, se necessário.

Isso é usado para normalizar as pontuações finais das sequências geradas e é opcional

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

Criar model_fn

Na prática, este será substituído por uma implementação modelo real, como aqui

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

Inicializar symbols_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

Ambicioso

Decodificação Greedy seleciona a ID de token com a maior probabilidade como seu próximo id: \(id_t = argmax_{w}P(id | id_{1:t-1})\) em cada iteração \(t\). O esboço a seguir mostra uma decodificação gananciosa.

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)

amostragem top_k

Na amostragem Top-K, o K provavelmente ids próximo token são filtrados e a massa de probabilidade é redistribuído entre apenas os ids 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)

amostragem top_p

Em vez de amostragem apenas do mais provável K símbolo ids, no Top-p amostragem escolhe entre o menor conjunto possível de ids cuja probabilidade cumulativa excede a probabilidade 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ção de pesquisa de feixe

A pesquisa de feixe reduz o risco de perder ids de token de alta probabilidade ocultos, mantendo o número de fluxos de hipóteses mais prováveis ​​em cada etapa de tempo e, eventualmente, escolhendo a hipótese que tem a maior probabilidade geral.

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)