{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!cp -r ../input/ai4code-source/* ./\n\nimport tensorflow as tf\nimport pandas as pd\nimport random\n\nfrom model.model import Model\nfrom data import Dataset\n\nTOKENIZERS_PARALLELISM = False","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:16:53.130545Z","iopub.execute_input":"2022-10-01T01:16:53.131052Z","iopub.status.idle":"2022-10-01T01:16:54.118622Z","shell.execute_reply.started":"2022-10-01T01:16:53.131006Z","shell.execute_reply":"2022-10-01T01:16:54.117266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Constant","metadata":{}},{"cell_type":"code","source":"MAX_LEN = 64\nMAX_CELL = 128\nBUFFER_SIZE = 32\nBATCH_SIZE = 2\nMODEL_PATH = 'microsoft/codebert-base'\nD_MODEL = 768\nNUM_TRAIN = 20000\nNUM_VAL = 5000\nDATA_DIR = \"../input/AI4Code\"\nCELL_PAD = 0\nLEARNING_RATE = 9e-7\nEPOCHS = 1\nDECAY_STEPS = 2000\nDECAY_RATE = 0.95","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:16:54.121217Z","iopub.execute_input":"2022-10-01T01:16:54.122122Z","iopub.status.idle":"2022-10-01T01:16:54.129264Z","shell.execute_reply.started":"2022-10-01T01:16:54.122060Z","shell.execute_reply":"2022-10-01T01:16:54.128052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data preparation","metadata":{}},{"cell_type":"markdown","source":"Import data","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv(\"../input/ai4code-split-dataset/data/train.csv\")\nval_df = pd.read_csv(\"../input/ai4code-split-dataset/data/val.csv\")\n# test = pd.read_csv(\"../input/ai4code-split-dataset/data/test.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:16:54.130546Z","iopub.execute_input":"2022-10-01T01:16:54.131548Z","iopub.status.idle":"2022-10-01T01:17:33.662811Z","shell.execute_reply.started":"2022-10-01T01:16:54.131508Z","shell.execute_reply":"2022-10-01T01:17:33.661783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Get a subset of notebooks for training","metadata":{}},{"cell_type":"code","source":"nb_ids = train_df.groupby('id').groups.keys()\n\nn = min(train_df.id.nunique(), NUM_TRAIN)\nrandom.seed(422)\nrandom_ids = random.sample(nb_ids, n)\n\ntrain_df = train_df[train_df.id.isin(random_ids)].reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:18:28.458983Z","iopub.execute_input":"2022-10-01T01:18:28.459663Z","iopub.status.idle":"2022-10-01T01:18:32.072421Z","shell.execute_reply.started":"2022-10-01T01:18:28.459628Z","shell.execute_reply":"2022-10-01T01:18:32.071420Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Optional. Get a small set for val set ==> make the validation step quicker\nnb_ids = val_df.groupby('id').groups.keys()\nn = min(val_df.id.nunique(), NUM_VAL)\nrandom.seed(42)\nrandom_ids = random.sample(nb_ids, n)\nval_df = val_df[val_df.id.isin(random_ids)].reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:18:32.079672Z","iopub.execute_input":"2022-10-01T01:18:32.079967Z","iopub.status.idle":"2022-10-01T01:18:32.289737Z","shell.execute_reply.started":"2022-10-01T01:18:32.079941Z","shell.execute_reply":"2022-10-01T01:18:32.288763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Sort dataset by ```pct_rank``` (**Only used for Phase 1**)","metadata":{}},{"cell_type":"code","source":"# train_df = train_df.sort_values(['id', 'pct_rank'])\n# val_df = val_df.sort_values(['id', 'pct_rank'])","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:18:32.294011Z","iopub.execute_input":"2022-10-01T01:18:32.294327Z","iopub.status.idle":"2022-10-01T01:18:32.300597Z","shell.execute_reply.started":"2022-10-01T01:18:32.294300Z","shell.execute_reply":"2022-10-01T01:18:32.299683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Create tensor dataset for training","metadata":{}},{"cell_type":"code","source":"d = Dataset(\n    model_path=MODEL_PATH, \n    max_len=MAX_LEN, \n    num_cells=MAX_CELL, \n    buffer_size=BUFFER_SIZE, \n    batch_size=BATCH_SIZE, \n    cell_pad=CELL_PAD\n)\n\n_, train = d.build_dataset(train_df)\n_, val = d.build_dataset(val_df)\n# test = d.build_dataset(test)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:18:32.303682Z","iopub.execute_input":"2022-10-01T01:18:32.303958Z","iopub.status.idle":"2022-10-01T01:20:13.482890Z","shell.execute_reply.started":"2022-10-01T01:18:32.303934Z","shell.execute_reply":"2022-10-01T01:20:13.481787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"# Input definitions\ninput_ids = tf.keras.layers.Input(shape=(MAX_CELL, MAX_LEN), name='input_ids', dtype='int32')\nattention_mask = tf.keras.layers.Input(shape=(MAX_CELL, MAX_LEN), name='attention_mask', dtype='int32')\ncell_features = tf.keras.layers.Input(shape=(MAX_CELL, 2), name='cell_features', dtype='float32')\ncell_mask = tf.keras.layers.Input(shape=(MAX_CELL), name='cell_mask', dtype='float32')\n\n# Model initialization\nstacked_transformer = Model(\n    model_path=MODEL_PATH, \n    d_model=D_MODEL, \n    n_heads=6, \n    num_cells=MAX_CELL,\n    dropout_pos=0.,\n    dropout_trans=0., \n    eps=1e-6, \n    d_ff_trans=64, \n    ff_activation=\"relu\", \n    n_layers=4,\n    d_ff_pointwise=64, \n    dropout_pointwise=0.\n)\n\n# Define the model\nmodel = tf.keras.Model(\n    inputs=[input_ids, attention_mask, cell_features, cell_mask], \n    outputs=stacked_transformer.call(input_ids, attention_mask, cell_features, cell_mask, True)\n)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:20:13.484886Z","iopub.execute_input":"2022-10-01T01:20:13.485584Z","iopub.status.idle":"2022-10-01T01:20:39.283934Z","shell.execute_reply.started":"2022-10-01T01:20:13.485546Z","shell.execute_reply":"2022-10-01T01:20:39.282961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.layers[2].trainable = False","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:20:39.285437Z","iopub.execute_input":"2022-10-01T01:20:39.285782Z","iopub.status.idle":"2022-10-01T01:20:39.299196Z","shell.execute_reply.started":"2022-10-01T01:20:39.285748Z","shell.execute_reply":"2022-10-01T01:20:39.298124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.load_weights(\"../input/ai4codemodel/saved_model/\")","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:20:39.302595Z","iopub.execute_input":"2022-10-01T01:20:39.303394Z","iopub.status.idle":"2022-10-01T01:20:43.190057Z","shell.execute_reply.started":"2022-10-01T01:20:39.303358Z","shell.execute_reply":"2022-10-01T01:20:43.189008Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Model summary\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:20:43.191574Z","iopub.execute_input":"2022-10-01T01:20:43.191932Z","iopub.status.idle":"2022-10-01T01:20:43.209246Z","shell.execute_reply.started":"2022-10-01T01:20:43.191897Z","shell.execute_reply":"2022-10-01T01:20:43.208028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n    initial_learning_rate=LEARNING_RATE,\n    decay_steps=DECAY_STEPS,\n    decay_rate=DECAY_RATE,\n    staircase=True\n)\noptimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)\n\nmodel.compile(optimizer=optimizer, loss=\"mse\", metrics=[\"mae\"])","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:20:43.210620Z","iopub.execute_input":"2022-10-01T01:20:43.211773Z","iopub.status.idle":"2022-10-01T01:20:43.233470Z","shell.execute_reply.started":"2022-10-01T01:20:43.211736Z","shell.execute_reply":"2022-10-01T01:20:43.232536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = model.fit(train, validation_data=val, epochs=EPOCHS)","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:20:43.234859Z","iopub.execute_input":"2022-10-01T01:20:43.235478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Save the Model","metadata":{}},{"cell_type":"code","source":"model.save_weights(\"saved_model/\")","metadata":{"execution":{"iopub.status.busy":"2022-10-01T01:17:36.743493Z","iopub.status.idle":"2022-10-01T01:17:36.744320Z","shell.execute_reply.started":"2022-10-01T01:17:36.744044Z","shell.execute_reply":"2022-10-01T01:17:36.744086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}