diciembre 28, 2021
~ 15 MIN
Sensio CoPilot
< Blog RSSSensio 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 🤗