diciembre 28, 2021

~ 15 MIN

Sensio CoPilot

< Blog RSS

Open In Colab

Sensio CoPilot

En este post vamos a entrenar una red neuronal para generación de código similar al funcionamiento de Github Copilot. Llevo usando Copilot unos meses y la verdad que puedo decir que es una herramienta increíble, que ha aumentado considerablemente mi productividad como programador. Hace unos días leí este hilo en Twitter y creí que sería interesante intentar replicar un sistema como Copilot desde cero. Así que sin más dilación, ¡vamos a ello!

Github Copilot utiliza el modelo conocido como Codex, desarrollado por OpenAI. Este modelo está basado en la arquitectura GPT (de lo que hablaremos más adelante) y tuneado con código público extraído de Github. Esto significa que el modelo fue pre-entrenado de manera no supervisada con mucho texto y luego se hizo fine tuning con el código extraído de Github para la tarea de generación de texto (nuevo código autocompletado). En contraste, nosotros entrenaremos un modelo desde cero para la tarea de generación de texto con un dataset preparado para ello a modo de demostración.

El Dataset

El dataset usado por OpenAI para entrenar Codex fue extraído de 54 millones de reposiotrios públicos de Github, conteniendo 179 GB de archivos Python de menos de 1 MB. Tras varias etapas de procesado, el dataset final ocupó 159 GB. Como no tenemos acceso a este dataset, usaremos CodeParrot, un dataset elaborado por HuggingFace para la tarea de generación de código.

Recuerda que OpenAI al final tiene que ganar dinero de alguna manera, y ésta es cobrando por el uso de sus modelos a través de la API. Siendo el modelo público, su única ventaja competitiva reside en los datos usados durante el entrenamiento. Esto es una tendencia clara en el mundo del Software 2.0, dónde el valor real está en los datos y no en el código (aunque como diría Andrej Karpathy en el Software 2.0 los datos SON el código y el modelo no es más que el binario resultante de la compilación del mismo, lo que llamamos el proceso de entrenamiento).

El dataset ocupa unos 50 GB, aproximadamente una tercera parte del dataset usado originalmente por OpenAI. ¡Nada mal! Puedes descargarlo utilizando los siguientes comandos:

git clone https://huggingface.co/datasets/lvwerra/codeparrot-clean-train
git clone https://huggingface.co/datasets/lvwerra/codeparrot-clean-valid

Para poder descargarlos necesitarás instalar Git LFS.

from pathlib import Path
from glob import glob

path = Path('data/codeparrot-clean-train')

files = glob(str(path) + '/*.json.gz')
len(files), files[:3]
(52,
 ['data/codeparrot-clean-train/file-000000000007.json.gz',
  'data/codeparrot-clean-train/file-000000000053.json.gz',
  'data/codeparrot-clean-train/file-000000000026.json.gz'])
import pandas as pd

file = pd.read_json(files[2], lines=True)
file

repo_name path copies size content license hash line_mean line_max alpha_frac autogenerated
0 jalavik/inspire-next setup.py 1 4558 # -*- coding: utf-8 -*-\n#\n# This file is par... gpl-2.0 -4849180608980663294 27.848101 77 0.615840 False
1 dlzhangxg/cloud-ml-sdk cloud_ml_samples/keras/mnist/trainer/task.py 1 2967 # Copyright 2017 Xiaomi, Inc.\n#\n# Licensed u... apache-2.0 -1822461891537938192 30.231579 74 0.649815 False
2 openstack/heat heat/engine/support.py 1 2683 #\n# Licensed under the Apache License, Ver... apache-2.0 -1815098437948811103 36.788732 78 0.622438 False
3 imapp-pl/golem tests/golem/network/test_golem_protocol.py 1 1444 import unittest\nfrom devp2p.service import Wi... gpl-3.0 -5671972220290336808 37.000000 77 0.654432 False
4 willimoa/pydal pydal/dialects/mongo.py 1 22083 import re\nfrom .._compat import PY2, basestri... bsd-3-clause 7148857323817256703 34.389423 93 0.518951 False
... ... ... ... ... ... ... ... ... ... ... ...
99995 georgestarcher/TA-SyncKVStore bin/ta_synckvstore/cloudconnectlib/core/ext.py 1 10300 import calendar\nimport json\nimport re\nimpor... mit -2116068727318744484 29.654762 80 0.593495 False
99996 nabucosound/django-propaganda propaganda/migrations/0001_initial.py 1 2828 # -*- coding: utf-8 -*-\nfrom __future__ impor... bsd-3-clause 2285667758777388556 38.830986 114 0.541372 False
99997 znuxor/aoc2016 4.py 1 41871 #!/usr/bin/python3\nimport operator\n\n# room ... bsd-3-clause -2104829268388077325 40.662687 118 0.834659 False
99998 sharadagarwal/autorest AutoRest/Generators/Python/Python.Tests/Expect... 1 5365 # coding=utf-8\n# ----------------------------... mit -2646046330546511516 35.250000 110 0.630009 False
99999 DeathSurvivorDE/dhbw_schreitbagger Schreitbagger/Bagger_GUI_v0-0-0-3_st.py 1 8009 \n'''\nBagger_GUI v0.0.0.1\n\nGrafische Benutz... gpl-3.0 54990144954214641 54.020690 212 0.507145 False

100000 rows × 11 columns

file.content[10]
'#!/usr/bin/env python\n\n# Copyright (C) 2014 Aldebaran Robotics\n#\n# Licensed under the Apache License, Version 2.0 (the "License");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an "AS IS" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n#import ROS dependencies\nimport rospy\n\n#import NAO dependencies\nfrom naoqi_driver.naoqi_node import NaoqiNode\nfrom geometry_msgs.msg import PoseStamped\nfrom geometry_msgs.msg import Pose\nimport almath\nimport tf\nfrom tf.transformations import euler_from_quaternion\n\nclass MoveToListener(NaoqiNode):\n\n    def __init__(self):\n        NaoqiNode.__init__(self, \'naoqi_moveto_listener\')\n        self.connectNaoQi()\n        self.listener = tf.TransformListener()\n\n        self.subscriber = rospy.Subscriber("/move_base_simple/goal", PoseStamped, self.callback)\n\n    # (re-) connect to NaoQI:\n    def connectNaoQi(self):\n        rospy.loginfo("Connecting to NaoQi at %s:%d", self.pip, self.pport)\n\n        self.motionProxy = self.get_proxy("ALMotion")\n        if self.motionProxy is None:\n            exit(1)\n\n    def callback(self, poseStamped):\n        # reset timestamp because of bug: https://github.com/ros/geometry/issues/82\n        poseStamped.header.stamp = rospy.Time(0)\n        try:\n            robotToTarget1 = self.listener.transformPose("/base_footprint", poseStamped)\n        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException) as e:\n            rospy.logerr("Error while transforming pose: %s", str(e))\n            return\n        quat = robotToTarget1.pose.orientation\n        (roll,pitch,yaw) = euler_from_quaternion((quat.x, quat.y, quat.z, quat.w))\n        self.motionProxy.moveTo(robotToTarget1.pose.position.x, robotToTarget1.pose.position.y, yaw)\n'

En posts anteriores en los que ya hemos hablado sobre aplicaciones de NLP vimos que a una red neuronal no podemos darle texto como entrada, sino números. El proceso de convertir las palabras de nuestro dataset en números se conoce como tokenización. Una opción sería construir un vector tan largo como número de palabras diferentes haya en el dataset y asignar un valor de 1 a la posición determinada por la palabra en concreto. Esta forma de tokenización se conoce como one-hot-encoding y como puedes imaginar es muy ineficiente. Lo más común consiste en convertir cada palabra en número entero y luego darle a la red como entrada la fila correspondiente en una matriz de embedding con dimensionalidad fijada por nosotros. Esto es mucho más eficiente, ya que el vector será denso y también permite establecer relaciones matemáticas entre palabras muy útil para la red. Si bien puedes implementar tu propia lógica de tokenización aquí usaremos una ya dado por la librería transformers, de esta manera nos evitaremos muchos dolores de cabeza como diferenciar entre palabras en mayúscula o minúscula, tratar los símbolos, etc. En este caso usaremos el tokenizador usado por el modelo BERT.

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
tokens = tokenizer.tokenize('def ADD_numbers(a, b):\n return ')

tokens
['def', 'add', '_', 'numbers', '(', 'a', ',', 'b', ')', ':', 'return']
indexes = tokenizer.convert_tokens_to_ids(tokens)

indexes
[13366, 5587, 1035, 3616, 1006, 1037, 1010, 1038, 1007, 1024, 2709]

Es importante tener en cuenta que el lenguaje usado en programación es mucho más restringido que el lenguaje natural, por lo que usar un tokenizer custom es probablemente una muy buena idea, reduciendo considerablemente el número de tokens necesarios y aumentando así la eficiencia de nuestro modelo.

Puedes aprender más sobre tokinizers aquí. Como verás existen muchas aproximaciones al problema, desde soluciones tan simples como contar el número de ocurrencias de cada palabra y usar su posición en la lista ordenada de todas las ocurrencias como token hasta reglas más complicadas. A medida que los transformers se hacen más grandes y potentes, pudiendo procesar longitdes más largas a la entrada, la última moda es usar directamente la representación ASCII de los caracteres, evitando cualquier tipo de tokenización 😂 (de lo cual soy muy fan).

x = 'Hola'

y = [ord(c) for c in x]
y
[72, 111, 108, 97]

El Modelo

El modelo usado por OpenAI, Codex, está basado en la arquitectura GPT y contiene 12 billones de parámetros. Esto está un poco fuera de nuestra alcance (de momento 😝) así que usaremos la implementación de Karpathy, minGPT, que nos permite entrenar pequeños transformers basados en la arquitectura GPT.

git clone https://github.com/karpathy/minGPT.git

Puedes jugar con el tamaños de la entrada (el número de token a usar), el número de capas, cabezas por capa y dimensión interna de las capas de atención. Nuestro modelo recibirá a la entrada un batch de trozos de código y nos dará a la salida un distribución de probabilidad sobre el vocabulario, dónde el valor más alto corresponderá a aquella palabra que el modelo cree que irá a continuación.

from minGPT.mingpt.model import GPT, GPTConfig

mconf = GPTConfig(tokenizer.vocab_size, block_size=512, n_layer=8, n_head=8, n_embd=512)
model = GPT(mconf)

model
GPT(
  (tok_emb): Embedding(30522, 512)
  (drop): Dropout(p=0.1, inplace=False)
  (blocks): Sequential(
    (0): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (1): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (2): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (3): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (4): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (5): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (6): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
    (7): Block(
      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (ln2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
      (attn): CausalSelfAttention(
        (key): Linear(in_features=512, out_features=512, bias=True)
        (query): Linear(in_features=512, out_features=512, bias=True)
        (value): Linear(in_features=512, out_features=512, bias=True)
        (attn_drop): Dropout(p=0.1, inplace=False)
        (resid_drop): Dropout(p=0.1, inplace=False)
        (proj): Linear(in_features=512, out_features=512, bias=True)
      )
      (mlp): Sequential(
        (0): Linear(in_features=512, out_features=2048, bias=True)
        (1): GELU()
        (2): Linear(in_features=2048, out_features=512, bias=True)
        (3): Dropout(p=0.1, inplace=False)
      )
    )
  )
  (ln_f): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
  (head): Linear(in_features=512, out_features=30522, bias=False)
)
import torch

x = torch.tensor([indexes]).long()
output, _ = model(x)
output.shape
torch.Size([1, 11, 30522])

GPT implementa un tipo de mecanismo de atención causal, esto significa que cada palabra solo puede atender a las palabras anteriores. De esta manera evitamos que el modelo sea capaz de ver en el futuro durante el entrenamiento. En fase de inferencia, tendremos que predecir palabra a palabra, usando las predicciones como nuevos inputs para seguir prediciendo.

from minGPT.mingpt.utils import sample

new_words = 3
y = sample(model, x, new_words, temperature=1.0, sample=True, top_k=10)[0]

y = tokenizer.convert_ids_to_tokens(y)
' '.join(y)
'def add _ numbers ( a , b ) : return invention shelf isaac'

Podemos quedarnos con aquellas palabras con la mayor probabilidad o, lo que es más común, samplear opciones siguiendo la distribución de probabilidad para generar diferentes posibilidades. Sin embargo, para que todo funciones, antes tendremos que entrenar nuestro modelo para que aprenda a generar código.

Entrenamiento

Como siempre os recomiendo, vamos a empezar haciendo el fit de una sola muestra para asegurarnos que todo está bien. Como los transformers trabajan con entradas de longitud fija vamos a añadir el token pad para rellenar los huecos que haga falta. Además, añadiremos los tokens de inicio y final de frase para que el modelo sepa cuando terminar de predecir nuevas palabras.

max_input_length = tokenizer.max_model_input_sizes['bert-base-uncased']

class Dataset(torch.utils.data.Dataset):
    def __init__(self, examples):
        self.examples = examples

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, ix):
        tokens = tokenizer.tokenize(self.examples[ix])
        tokens = tokens[:max_input_length-2]

        input = [tokenizer.cls_token_id] + \
            tokenizer.convert_tokens_to_ids(tokens) + \
            [tokenizer.pad_token_id] * (512 - len(tokens) - 1)

        target = tokenizer.convert_tokens_to_ids(tokens) + \
            [tokenizer.sep_token_id] + \
            [tokenizer.pad_token_id] * (512 - len(tokens) - 1)

        assert len(input) == 512
        assert len(target) == 512

        return torch.LongTensor(input), torch.LongTensor(target)
ds = Dataset(['def add_numbers(a, b):\n return a + b'])

inputs, target = ds[0]
inputs[:17], target[:17]
(tensor([  101, 13366,  5587,  1035,  3616,  1006,  1037,  1010,  1038,  1007,
          1024,  2709,  1037,  1009,  1038,     0,     0]),
 tensor([13366,  5587,  1035,  3616,  1006,  1037,  1010,  1038,  1007,  1024,
          2709,  1037,  1009,  1038,   102,     0,     0]))

El target es igual que el input pero corrido un token a la derecha, así para cada input el modelo intentará predecir la siguiente palabra. El primer token del input es el token de inicio de frase, para lo que el modelo nos dará la primera palabra. El último token del target es el token de final de frase, así cuando el modelo reciba la última palabra nos dará este token indicando que ya no debe predecir más. El token de padding sirve simplemente para tener una frase de longitud determinada.

dl = torch.utils.data.DataLoader(ds, batch_size=1, shuffle=False)

x, y = next(iter(dl))
x.shape, y.shape
(torch.Size([1, 512]), torch.Size([1, 512]))

Para entrenar, usaremos Pytorch Lightning.

import pytorch_lightning as pl

class Module(pl.LightningModule):
    def __init__(self, model, tconf):
        super().__init__()
        self.model = model
        self.tconf = tconf

    def forward(self, x):
        return self.model(x)[0]

    def training_step(self, batch, batch_idx):
        y, loss = self.model(*batch)
        return loss

    def validation_step(self, batch, batch_idx):
        y, loss = self.model(*batch)
        self.log('val_loss', loss.item())

    def configure_optimizers(self):
        return self.model.configure_optimizers(self.tconf)
from minGPT.mingpt.trainer import TrainerConfig

mconf = GPTConfig(tokenizer.vocab_size, block_size=512, n_layer=8, n_head=8, n_embd=512)
model = GPT(mconf)
tconf = TrainerConfig(weight_decay=0.1, learning_rate=3e-4, betas=(0.9, 0.95))

module = Module(model, tconf)
trainer = pl.Trainer(
    gpus=1,
    max_epochs=100,
    precision=16,
    overfit_batches=1,
    logger=False,
    checkpoint_callback=False,
)

trainer.fit(module, dl)
Using 16bit native Automatic Mixed Precision (AMP)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name  | Type | Params
-------------------------------
0 | model | GPT  | 56.7 M
-------------------------------
56.7 M    Trainable params
0         Non-trainable params
56.7 M    Total params
113.474   Total estimated model params size (MB)
/home/juan/miniconda3/lib/python3.9/site-packages/pytorch_lightning/trainer/data_loading.py:111: UserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 20 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(



Training: 0it [00:00, ?it/s]
tokens = tokenizer.tokenize('def add_numbers(a, b):')
indexes = [tokenizer.cls_token_id] + tokenizer.convert_tokens_to_ids(tokens)

x = torch.tensor([indexes]).long()
new_words = 5
y = sample(model, x, new_words)[0]

y = tokenizer.convert_ids_to_tokens(y)
' '.join(y)
'[CLS] def add _ numbers ( a , b ) : return a + b [SEP]'

!Genial! Nuestro modelo ha sido capaz de aprenderse de memoria una sola muestra. Ahora tenemos que entrenar con todo el dataset.

train_path = Path('data/codeparrot-clean-train')
val_path = Path('data/codeparrot-clean-valid')

train_files = glob(str(train_path) + '/*.json.gz')
val_files = glob(str(val_path) + '/*.json.gz')

len(train_files), len(val_files)
(52, 1)
train_file = pd.read_json(train_files[0], lines=True) # le paso sólo un archivo, no se cómo hacer que los use todos 🥲
val_file = pd.read_json(val_files[0], lines=True)
ds = {
    'train': Dataset(train_file.content.values),
    'val': Dataset(val_file.content.values),
}

len(ds['train']), len(ds['val'])
(100000, 61373)
dl = {
    'train': torch.utils.data.DataLoader(ds['train'], batch_size=32, shuffle=True, num_workers=20, pin_memory=True),
    'val': torch.utils.data.DataLoader(ds['val'], batch_size=32, shuffle=False, num_workers=20, pin_memory=True),
}
mconf = GPTConfig(tokenizer.vocab_size, block_size=512, n_layer=8, n_head=8, n_embd=512)
model = GPT(mconf)
tconf = TrainerConfig(weight_decay=0.1, learning_rate=3e-4, betas=(0.9, 0.95))

module = Module(model, tconf)
trainer = pl.Trainer(
    gpus=1,
    max_epochs=3,
    precision=16,
    limit_val_batches=100,
)

trainer.fit(module, dl['train'], dl['val'])
Using 16bit native Automatic Mixed Precision (AMP)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name  | Type | Params
-------------------------------
0 | model | GPT  | 56.7 M
-------------------------------
56.7 M    Trainable params
0         Non-trainable params
56.7 M    Total params
113.474   Total estimated model params size (MB)



Validation sanity check: 0it [00:00, ?it/s]



Training: 0it [00:00, ?it/s]



Validating: 0it [00:00, ?it/s]



Validating: 0it [00:00, ?it/s]



Validating: 0it [00:00, ?it/s]
tokens = tokenizer.tokenize('def add_numbers(a, b):')
indexes = [tokenizer.cls_token_id] + tokenizer.convert_tokens_to_ids(tokens)

x = torch.tensor([indexes]).long()
new_words = 5
y = sample(model, x, new_words)[0]

y = tokenizer.convert_ids_to_tokens(y)
' '.join(y)
'[CLS] def add _ numbers ( a , b ) : return a + b [SEP]'

¡Voilá! Ya tenemos nuestro modelo GPT para generación de cógido 🔥 Para obtener un buen resultado tendrás que entrenar más epochs, con todos los datos, un modelo más grande (puedes usar los papers de GPT2 y GPT3 para guiarte), etc.

Resumen

En este post hemos visto cómo entrenar una red neuronal para generación de código Python inspirándonos en Github Copilot. Aquí he utilizado diferentes librerías y herramientas, pero como comentaba al principio del post en este hilo encontrarás cómo hacer lo mismo con las librerías de Huggingface, lo cuál simplifica mucho todo y probablemete te permita conseguir muchos mejores resultados. En próximos posts entraremos más en detalle en el ecosistema de Huggingface. Te animo a que juegues un poco con el código, entrenando diferentes versiones del modelo. Para pasar al siguiente nivel, puedes exportar el modelo y desplegarlo en una API, hosteada en Heroku por ejemplo, y hacer una extensión de VSCode que llame a la API para sugerir código directamente en el editor. Si lo haces, compártelo en nuestro discord con el resto de la comunidad 🤗

< Blog RSS