Cómo implementar redes Transformer en Python para traducción Inglés a Español

Las Redes Transformer han generado una verdadera revolución en el área del Procesamiento de Lenguaje Natural (NLP). A diferencia de las redes recurrentes tradicionales, los Transformers permiten procesar secuencias completas de texto de forma simultánea, optimizando el entrenamiento y capturando relaciones complejas entre palabras sin importar su distancia en la frase.

En este tutorial, veremos paso a paso cómo programar un traductor de inglés a español en Python, utilizando la librería keras-transformer en Google Colab. Con esto tendremos un traductor capaz de recibir una frase en Inglés y generar su traducción correspondiente al español.

Al final del tutorial se encuentra el enlace para descargar el código fuente.

Introducción a la Arquitectura Transformer

Antes de entrar al código, es fundamental recordar la estructura básica de un una Red Transformer. Este modelo se basa en dos grandes estructuras:

El Codificador (Encoder): Procesa la frase original (en este caso, en Inglés). El proceso comienza con un embedding de entrada donde cada palabra se convierte en un vector. Dado que el Transformer procesa todo simultáneamente y no palabra por palabra, se aplica una Codificación Posicional para que el modelo entienda el orden del texto. Los bloques del codificador aprenden representaciones numéricas (embeddings) de la frase prestando atención a los elementos circundantes usando los bloques de Atención (Self-Attention). En esencia, estos bloques de atención son los que permiten al modelo «entender» el texto.
El Decodificador (Decoder): Es el encargado de generar la frase en español. A diferencia del codificador, el decodificador genera el texto palabra por palabra. Además, utiliza un token de inicio (<START>) y se apoya tanto en la información procesada por el codificador como en los tokens ya generados para predecir la siguiente palabra mediante una capa de salida que usa una función de activación Softmax.

Para agilizar el desarrollo usaremos la librería keras-transformer, que nos permite implementar estas capas complejas sin necesidad de escribir demasiadas líneas de código.

Configuración del Entorno y Datos

Para este proyecto, utilizaremos un set de datos con aproximadamente 125,000 frases en inglés con su equivalente en español. Y además, usaremos Google Colab para aprovechar la GPU que nos proporciona esta plataforma.

Instalación de Librerías

El primer paso es instalar la librería necesaria para manejar los bloques atencionales y la estructura del Transformer.

Para ello simplemente ejecutamos la siguiente línea de código directamente desde el notebook de Python:

!pip install keras-transformer

Importación de Módulos

Importaremos numpy para la manipulación de arreglos, pickle para la lectura de los datos pre-procesados y los módulos específicos de keras-transformer:

import numpy as np
from keras_transformer import get_model, decode
from pickle import load
from google.colab import drive

# Establecemos una semilla para garantizar la reproducibilidad del entrenamiento
np.random.seed(0)

Lectura del Dataset

Los datos se encuentran almacenados en Google Drive en formato .pkl.

Así que primero montamos Google Drive en nuestra máquina de Google Colab y luego cargamos el set de entrenamiento:

# Montar Google Drive
drive.mount('/content/drive')

# Definir la ruta y cargar el dataset
filename = '/content/drive/My Drive/english-spanish.pkl'
dataset = load(open(filename, 'rb'))

# Imprimimos un ejemplo (el dato 120,000) para verificar
print(dataset[120000,0]) # Frase en inglés
print(dataset[120000,1]) # Frase en español

Tras la carga (línea 6 del bloque de código anterior), imprimimos en pantalla un dato de ejemplo (líneas 9 y 10). Al hacer esto observamos algo como esto:

tom is a new yorker but he doesnt have a new york accent
tom es neoyorquino pero no tiene acento de nueva york

Es decir que cada dato contiene simplemente un par de frases Inglés-Español.

Preparación del dataset

Tokenización

El primer paso en este pre-procesamiento de los datos es la tokenización, que consiste en separar cada frase en palabras individuales dentro de una lista. Para lograr esto podemos simplemente podemos usar algunos métodos avanzados para el procesamiento de strings en Python (como por ejemplo «split()») junto con herramientas para la manipulación de listas:

# Crear tokens para el idioma fuente (inglés)
source_tokens = []
for sentence in dataset[:,0]:
  source_tokens.append(sentence.split(' '))

# Crear tokens para el idioma objetivo (español)
target_tokens = []
for sentence in dataset[:,1]:
  target_tokens.append(sentence.split(' '))

# Verificamos los tokens del ejemplo anterior
print(source_tokens[120000])
print(target_tokens[120000])

En el código anterior simplemente tomamos cada frase en Inglés y con el uso de «split()» la dividimos en palabras (líneas 1 a 4) y hacemos algo similar para las frases en español (líneas 6 a 9). Al imprimir en pantalla algunos ejemplos de los tokens obtenidos (líneas 12 y 13) obtendremos algo similar a esto:

['tom', 'is', 'a', 'new', 'yorker', 'but', 'he', 'doesnt', 'have', 'a', 'new', 'york', 'accent']
['tom', 'es', 'neoyorquino', 'pero', 'no', 'tiene', 'acento', 'de', 'nueva', 'york']

Es decir, simplemente tendremos las mismas frases originales pero en listas de Python, donde cada palabra será precisamente un token.

Construcción de Diccionarios de Equivalencia

La Red Transformer, como cualquier tipo de Red Neuronal, procesa los datos en formato numérico. Así que debemos convertir los tokens obtenidos en el paso anterior en representaciones numéricas.

Para ello crearemos diccionarios de Python donde cada palabra única tendrá un número asignado. Además, definiremos tokens especiales que resultarán claves para que el modelo pueda interpretar correctamente el texto de entrada (en Inglés) y generar adecuadamente el texto de salida (en Español):

<PAD> (0): para rellenar frases cortas y que todas tengan la misma longitud.
<START> (1): indica el inicio de una secuencia.
<END> (2): indica el final de una secuencia.

Veamos esta implementación:

def build_token_dict(token_list):
  token_dict = {
      '<PAD>': 0,
      '<START>': 1,
      '<END>': 2
  }
  for tokens in token_list:
    for token in tokens:
      if token not in token_dict:
        token_dict[token] = len(token_dict)
  return token_dict

# Construimos los diccionarios
source_token_dict = build_token_dict(source_tokens)
target_token_dict = build_token_dict(target_tokens)

# Creamos un diccionario inverso para convertir números de vuelta a palabras en español
target_token_dict_inv = {v:k for k,v in target_token_dict.items()}

Desglosemos lo que hace esta función «build_token_dict()»:

En las líneas 2 a 6 incluimos las equivalencias para los tokens especiales
Mientras que en las líneas de la 7 a la 10 iteramos por cada token en la lista y si no está en el diccionario lo incluimos (línea 10). El número asignado será simplemente el tamaño del diccionario, el cual cambiará dinámicamente a medida que agregamos más y más tokens al mismo
Y como resultado retornamos el diccionario de equivalencias (línea 11)

En las líneas 14 y 15 simplemente llamamos la función anterior, lo cual nos permite crear los diccionarios para el texto origen en Inglés («source_token_dict») y para el texto destino en Español («target_token_dict»).

Además, como la Red Transformer generará a la salida tokens en formato numérico, es necesario crear un diccionario que tome dichas cantidades numéricas y las convierta al formato texto, con lo cual podremos ver el resultado de la traducción. Para ello creamos el diccionario «target_token_dic_inv» (línea 18) que permita hacer esta conversión.

Agregando Tokens de Control y Padding

Ahora debemos aplicar los tokens especiales para indicar el inicio y finalización de cada frase. Y además, debemos estandarizar sus longitudes usando el padding (o rellenado con ceros).

Este padding es necesario pues tanto el codificador como el decodificador de la Red Transformer aceptan a la entrada secuencias de tamaño fijo. Así que el padding permitirá simplemente completar con ceros aquellas secuencias que tengan relativamente pocas palabras.

Comencemos entonces añadiendo los tokens de inicio («<START>») y finalización («<END>») a cada secuencia:

# Agregar tokens de inicio y fin
encoder_tokens = [['<START>'] + tokens + ['<END>'] for tokens in source_tokens]
decoder_tokens = [['<START>'] + tokens + ['<END>'] for tokens in target_tokens]
output_tokens = [tokens + ['<END>'] for tokens in target_tokens]

Y para realizar el padding simplemente calcularemos la longitud máxima de las frases en ambos idiomas para rellenar las más cortas con el token <PAD>:

# Determinar longitudes máximas
source_max_len = max(map(len, encoder_tokens))
target_max_len = max(map(len, decoder_tokens))

# Aplicar Padding
encoder_tokens = [tokens + ['<PAD>']*(source_max_len-len(tokens)) for tokens in encoder_tokens]
decoder_tokens = [tokens + ['<PAD>']*(target_max_len-len(tokens)) for tokens in decoder_tokens]
output_tokens = [tokens + ['<PAD>']*(target_max_len-len(tokens)) for tokens in output_tokens ]

Por ejemplo, si ahora imprimimos el dato número 120.000 del encoder, que corresponde a la frase en Inglés «tom is a new yorker but he doesn’t have a new york accent», veremos que la secuencia resultante será:

['<START>', 'tom', 'is', 'a', 'new', 'yorker', 'but', 'he', 'doesnt', 'have', 'a', 'new', 'york', 'accent', '<END>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>', '<PAD>']

Con lo cual podemos comprobar que los tokens de la frase en mención están delimitados por «<START>» y «<END>» y que, por tratarse de una secuencia relativamente corta, se ha rellenado con ceros en la parte final (tokens «<PAD>»).

Codificación Numérica de las Secuencias

Finalmente, transformamos las listas de palabras en listas de números usando los diccionarios creados anteriormente:

encoder_input = [list(map(lambda x: source_token_dict[x], tokens)) for tokens in encoder_tokens]
decoder_input = [list(map(lambda x: target_token_dict[x], tokens)) for tokens in decoder_tokens]
output_decoded = [list(map(lambda x: [target_token_dict[x]], tokens)) for tokens in output_tokens]

Así que en este punto ya tendremos codificada numéricamente cada una de las frases en el set de datos. Si por ejemplo imprimimos en pantalla el dato 120.000 codificado (variable «encoder_input») tendremos este resultado:

[1, 56, 258, 120, 197, 12666, 2914, 32, 1577, 140, 120, 197, 5385, 4287, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Y este es precisamente un ejemplo de las secuencias que ingresarán al modelo durante el entrenamiento.

¡Perfecto, y con todo esto ya estamos listos para construir nuestro traductor de texto Inglés a Español basado en Redes Transformer!

Implementación de la Red Transformer

Para crear una instancia de la Red Transformer usaremos la función «get_model» de la librería keras_transformer, importada al inicio del tutorial.

Al llamar esta función definiremos los siguientes argumentos que se ajustarán a nuestros recursos computacionales:

token_num: el tamaño máximo del vocabulario entre ambos diccionarios (Inglés y Español)
embed_dim: 32 (dimensión de los vectores de embedding).
encoder_num / decoder_num: 2 (número de bloques en cada estructura).
head_num: 4 (número de cabezales multi-atención).
hidden_dim: 128 (neuronas en la capa oculta de la red feed-forward).
dropout_rate: 0.05 (para evitar el sobreajuste o overfitting).

Este es el código que nos permite crear la instancia del modelo:

model = get_model(
    token_num = max(len(source_token_dict),len(target_token_dict)),
    embed_dim = 32,
    encoder_num = 2,
    decoder_num = 2,
    head_num = 4,
    hidden_dim = 128,
    dropout_rate = 0.05,
    use_same_embed = False,
)

model.compile('adam', 'sparse_categorical_crossentropy')
model.summary()

En la línea 12 del bloque de código anterior simplemente estamos definiendo el optimizador que usaremos para entrenar el modelo (Adam) así como la pérdida que guiará este proceso de entrenamiento.

Por otra parte, en la línea 13 usamos el método «summary()» para imprimir en pantalla las características de este modelo. Tras hacer esto podremos verificar que el modelo cuenta con aproximadamente 1.7 millones de parámetros entrenables. Es decir, se trata de un modelo relativamente pequeño para los tamaños usuales de este tipo de Redes (cuyo número de parámetros es del orden de varios cientos o miles de millones de parámetros).

Entrenamiento del modelo

El entrenamiento con 125,000 frases y 15 iteraciones puede ser lento (aprox. 10 minutos por iteración en GPU). Este sería el código requerido para este entrenamiento:

# Preparación de datos para Keras
x = [np.array(encoder_input), np.array(decoder_input)]
y = np.array(output_decoded)

# Entrenamiento
model.fit(x,y, epochs=15, batch_size=32)

En el bloque de código anterior:

En las líneas 2 y 3 definimos los arreglos de entrada («x») y salida («y») del modelo (resultado del pre-procesamiento que realizamos anteriormente)
Mientras que en la línea 6 usamos el método «fit» para entrenar el modelo por un total de 15 iteraciones

Traducción de texto con el modelo entrenado

Así que en este punto ya tenemos el modelo entrenado y ya estamos listos para ponerlo a prueba.

Así que crearemos la función translate, que toma una frase en inglés, la procesa, la convierte a números, la lleva al modelo que acabamos de entrenar y al final decodifica el resultado numérico generado por el modelo para obtener el texto en el lenguaje Español.

Esta es la implementación correspondiente:

def translate(sentence):
    # Tokenización y Padding de la frase de entrada
    sentence_tokens = [tokens + ['<END>', '<PAD>'] for tokens in [sentence.split(' ')]]
    tr_input = [list(map(lambda x: source_token_dict[x], tokens)) for tokens in sentence_tokens][0]
    
    # Decodificación usando el Transformer
    decoded = decode(
        model, 
        tr_input, 
        start_token = target_token_dict['<START>'],
        end_token = target_token_dict['<END>'],
        pad_token = target_token_dict['<PAD>']
    )

    # Convertir números a palabras (omitimos el token START en la impresión)
    translation = ' '.join(map(lambda x: target_token_dict_inv[x], decoded[1:-1]))
    
    print(f'Frase original (Inglés): {sentence}')
    print(f'Traducción (Español): {translation}')

En el código anterior:

En las líneas 3 y 4 realizamos el pre-procesamiento de la secuencia de entrada (tokenización, padding, tokens de inicio y finalización y codificación numérica)
En las líneas de la 7 a la 13 hacemos uso del método «decode» de la librería keras_transformer para presentar la secuencia de entrada codificada al modelo y generar la predicción en formato numérico
Finalmente, en la línea 16 tomamos esta secuencia numérica y usamos el diccionario «target_token_dict_inv» para representar esta secuencia de salida en formato de texto y finalmente imprimir tanto la secuencia de entrada como la de salida (líneas 18 y 19)

Usando esta función, podemos poner a prueba el el modelo con frases de distinta complejidad:

Patrones gramaticales: en inglés el adjetivo va antes del sujeto (red car), mientras que en español es al revés (auto rojo). Así que si para la frase «the red car» el modelo realiza la traducción correcatamente («el auto rojo»)
Tiempos verbales: cuando introducimos la frase «the cat eats the food» el modelo genera a la salida «el gato se comió la comida», es decir que erróneamente cambia el tiempo verbal al realizar la traducción.
Traducción intuitiva (no 1 a 1): pero si introducimos una frase más compleja al modelo, como por ejemplo «do you understand what I say?» el modelo demuestra entender que el número de palabras de entrada no debe ser igual al de salida, generando una traducción adecuada para esa frase de entrada («entiende lo que digo?»)

Enlaces de descarga set de datos y código fuente

▼ El código fuente de este tutorial

▼ El set de datos de este tutorial

Conclusiones

Aunque este modelo es bastante simple (pues tiene sólo unos pocos millones de parámetros), funciona sorprendentemente bien para las frases que probamos anteriormente.

Sin embargo, aún tiene varias limitaciones. Por ejemplo:

A veces confunde el género (mis favoritos en lugar de mis favoritas).
Frases muy largas o con vocabulario no presente en el entrenamiento pueden generar errores semánticos.
Un traductor profesional (como el de Google) utiliza muchas más capas, cabezales de atención y un set de datos masivo.

Sin embargo, este sencillo tutorial nos ha permitido ver el gran potencial que tienen las Redes Transformer en el Procesamiento del Lenguaje Natural. Si usamos un modelo más complejo (es decir con muchos más parámetros) así como un set de entrenamiento con muchos más datos, es posible construir un traductor de texto mucho más robusto y mucho más preciso.