AI4Code Pytorch DistilBert Baseline
导库
import json
from pathlib import Path
import numpy as np
import pandas as pd
from scipy import sparse
from tqdm import tqdm
pd.options.display.width = 180
pd.options.display.max_colwidth = 120
BERT_PATH = "../input/huggingface-bert-variants/distilbert-base-uncased/distilbert-base-uncased"
data_dir = Path('../input/AI4Code')
读取文件
NUM_TRAIN = 10000
def read_notebook(path):
return (
pd.read_json(
path,
dtype={'cell_type': 'category', 'source': 'str'})
.assign(id=path.stem)
.rename_axis('cell_id')
)
paths_train = list((data_dir / 'train').glob('*.json'))#[:NUM_TRAIN]
notebooks_train = [
read_notebook(path) for path in tqdm(paths_train, desc='Train NBs')
]
df = (
pd.concat(notebooks_train)
.set_index('id', append=True)
.swaplevel()
.sort_index(level='id', sort_remaining=False)
)
df
读取顺序文件
df_orders = pd.read_csv(
data_dir / 'train_orders.csv',
index_col='id',
squeeze=True,
).str.split() # Split the string representation of cell_ids into a list
df_orders
示例文件展示
# Get an example notebook
nb_id = df.index.unique('id')[6]
print('Notebook:', nb_id)
print("The disordered notebook:")
nb = df.loc[nb_id, :]
display(nb)
print()
cell_order = df_orders.loc[nb_id]
print("The ordered notebook:")
nb.loc[cell_order, :]
无序的:
有序的:
获取df_ranks
def get_ranks(base, derived):
"""获取每个cell_id应该放到对应位置的索引"""
return [base.index(d) for d in derived]
cell_ranks = get_ranks(cell_order, list(nb.index))
nb.insert(0, 'rank', cell_ranks)
nb
根据id获取有序cell_id和无序的cell_id
df_orders_ = df_orders.to_frame().join(
df.reset_index('cell_id').groupby('id')['cell_id'].apply(list),
how='right',
)
df_orders_
ranks = {}
for id_, cell_order, cell_id in df_orders_.itertuples():
ranks[id_] = {'cell_id': cell_id, 'rank': get_ranks(cell_order, cell_id)}
df_ranks = (
pd.DataFrame
.from_dict(ranks, orient='index')
.rename_axis('id') # 给两个轴[index,columns]设置名称
.apply(pd.Series.explode) # 将列表拆分成一行行的内容
.set_index('cell_id', append=True) # 将每个id的内容拼接起来
)
df_ranks
获取祖先信息
df_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')
df_ancestors
将祖先信息、ranks信息到df
df = df.reset_index().merge(df_ranks, on=["id", "cell_id"]).merge(df_ancestors, on=["id"])
df
df["pct_rank"] = df["rank"] / df.groupby("id")["cell_id"].transform("count") # 计算每个notebook有多少cell,然后对rank进行归一化处理 其实是定义好了cell的排列顺序
# 对归一化后的rank进行分箱统计
df["pct_rank"].hist(bins=10)
数据预处理
划分训练集和验证集
from sklearn.model_selection import GroupShuffleSplit
NVALID = 0.1 # size of validation set
splitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)
train_ind, val_ind = next(splitter.split(df, groups=df["ancestor_id"])) # 按照组别划分 详细可见reference 2
train_df = df.loc[train_ind].reset_index(drop=True)
val_df = df.loc[val_ind].reset_index(drop=True)
train_df.head()
针对不同cell_type处理数据
train_df_mark = train_df[train_df["cell_type"] == "markdown"].reset_index(drop=True)
val_df_mark = val_df[val_df["cell_type"] == "markdown"].reset_index(drop=True)
衡量指标计算函数
from bisect import bisect
def count_inversions(a):
inversions = 0
sorted_so_far = []
for i, u in enumerate(a):
j = bisect(sorted_so_far, u)
inversions += i - j
sorted_so_far.insert(j, u)
return inversions
def kendall_tau(ground_truth, predictions):
total_inversions = 0
total_2max = 0 # twice the maximum possible inversions across all instances
for gt, pred in zip(ground_truth, predictions):
ranks = [gt.index(x) for x in pred] # rank predicted order in terms of ground truth
total_inversions += count_inversions(ranks)
n = len(gt)
total_2max += n * (n - 1)
return 1 - 4 * total_inversions / total_2max
计算示例:
y_dummy = val_df.groupby('id')['cell_id'].apply(list)
kendall_tau(df_orders.loc[y_dummy.index], y_dummy)
0.42511216883092573
数据建模
选择NLP模型
from tqdm import tqdm
import sys, os
from transformers import DistilBertModel, DistilBertTokenizer
import torch.nn.functional as F
import torch.nn as nn
import torch
MAX_LEN = 128
class MarkdownModel(nn.Module):
def __init__(self):
super(MarkdownModel, self).__init__()
self.distill_bert = DistilBertModel.from_pretrained(BERT_PATH)
self.top = nn.Linear(768, 1)
def forward(self, ids, mask):
x = self.distill_bert(ids, mask)[0]
# 使用最后一层CLS做投影到1维数值
x = self.top(x[:, 0, :])
return x
数据格式封装
from torch.utils.data import DataLoader, Dataset
class MarkdownDataset(Dataset):
def __init__(self, df, max_len):
super().__init__()
self.df = df.reset_index(drop=True)
self.max_len = max_len
self.tokenizer = DistilBertTokenizer.from_pretrained(BERT_PATH, do_lower_case=True)
def __getitem__(self, index):
row = self.df.iloc[index]
inputs = self.tokenizer.encode_plus(
row.source,
None,
add_special_tokens=True,
max_length=self.max_len,
padding="max_length",
return_token_type_ids=True,
truncation=True
)
ids = torch.LongTensor(inputs['input_ids'])
mask = torch.LongTensor(inputs['attention_mask'])
return ids, mask, torch.FloatTensor([row.pct_rank])
def __len__(self):
return self.df.shape[0]
train_ds = MarkdownDataset(train_df_mark, max_len=MAX_LEN)
val_ds = MarkdownDataset(val_df_mark, max_len=MAX_LEN)
val_ds[0]
BS = 32 # 设置batch_size
NW = 8 # 设置数据处理时,num_workers,表示使用多少个cpu核心来准备数据
train_loader = DataLoader(train_ds, batch_size=BS, shuffle=True, num_workers=NW,
pin_memory=False, drop_last=True)
val_loader = DataLoader(val_ds, batch_size=BS, shuffle=False, num_workers=NW,
pin_memory=False, drop_last=False)
设置工具函数
def adjust_lr(optimizer, epoch):
"""根据epoch调整学习率"""
if epoch < 1:
lr = 5e-5
elif epoch < 2:
lr = 1e-3
elif epoch < 5:
lr = 1e-4
else:
lr = 1e-5
for p in optimizer.param_groups:
p['lr'] = lr
return lr
def get_optimizer(net):
"""传入网络,返回优化器"""
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=3e-4, betas=(0.9, 0.999),
eps=1e-08)
return optimizer
def read_data(data):
"""将数据送入GPU,并将前n-1个数据按照元组返回,最后一个元素单独返回,意味着将X,y分开返回"""
return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()
定义训练、验证函数
def validate(model, val_loader):
model.eval()
tbar = tqdm(val_loader, file=sys.stdout)
preds = []
labels = []
with torch.no_grad():
for idx, data in enumerate(tbar):
inputs, target = read_data(data)
pred = model(inputs[0], inputs[1])
preds.append(pred.detach().cpu().numpy().ravel())
labels.append(target.detach().cpu().numpy().ravel())
return np.concatenate(labels), np.concatenate(preds) # 将每个batch的结果拼接在一起
def train(model, train_loader, val_loader, epochs):
np.random.seed(0)
optimizer = get_optimizer(model)
criterion = torch.nn.MSELoss()
for e in range(epochs):
model.train()
tbar = tqdm(train_loader, file=sys.stdout)
lr = adjust_lr(optimizer, e)
loss_list = []
preds = []
labels = []
for idx, data in enumerate(tbar):
inputs, target = read_data(data)
optimizer.zero_grad()
pred = model(inputs[0], inputs[1])
loss = criterion(pred, target)
loss.backward()
optimizer.step()
loss_list.append(loss.detach().cpu().item())
preds.append(pred.detach().cpu().numpy().ravel())
labels.append(target.detach().cpu().numpy().ravel())
avg_loss = np.round(np.mean(loss_list), 4)
tbar.set_description(f"Epoch {e+1} Loss: {avg_loss} lr: {lr}")
y_val, y_pred = validate(model, val_loader)
print("Validation MSE:", np.round(mean_squared_error(y_val, y_pred), 4))
print()
return model, y_pred
训练并预测
model = MarkdownModel() # 该模型可以预测数据集中markdown块的百分比顺序
model = model.cuda()
model, y_pred = train(model, train_loader, val_loader, epochs=3) # y_pred表示验证集的预测输出
val_df["pred"] = val_df.groupby(["id", "cell_type"])["rank"].rank(pct=True)
val_df
# 将markdown块的预测值 插入到pred 对于code块的pct_rank不做改变
val_df.loc[val_df["cell_type"] == "markdown", "pred"] = y_pred
val_df
# 将pred列的值进行排序 那么就可以得到每个id对应的cell_id顺序
y_dummy = val_df.sort_values("pred").groupby('id')['cell_id'].apply(list)
y_dummy
# 计算预测效果
kendall_tau(df_orders.loc[y_dummy.index], y_dummy)
# 0.6471850137972645
预测测试数据
读取测试集文件
paths_test = list((data_dir / 'test').glob('*.json'))
notebooks_test = [
read_notebook(path) for path in tqdm(paths_test, desc='Test NBs')
]
test_df = (
pd.concat(notebooks_test)
.set_index('id', append=True)
.swaplevel()
.sort_index(level='id', sort_remaining=False)
).reset_index()
test_df
# 测试集缺少祖先信息和rank信息
# 每个id初始化的rank设置为原始序号
test_df["rank"] = test_df.groupby(["id", "cell_type"]).cumcount() # 根据id cell_type进行累计计数
test_df
# 根据rank计算出pct_rank,并赋值给pred列
test_df["pred"] = test_df.groupby(["id", "cell_type"])["rank"].rank(pct=True)
test_df
预处理测试集
# 设置pct_rank初始化为0
test_df["pct_rank"] = 0
test_ds = MarkdownDataset(test_df[test_df["cell_type"] == "markdown"].reset_index(drop=True), max_len=MAX_LEN)
test_loader = DataLoader(test_ds, batch_size=BS, shuffle=False, num_workers=NW,
pin_memory=False, drop_last=False)
len(test_ds), test_ds[0]
_, y_test = validate(model, test_loader)
# 预测出markdown块的值
test_df.loc[test_df["cell_type"] == "markdown", "pred"] = y_test
保存文件
# 对prd进行排序
sub_df = test_df.sort_values("pred").groupby("id")["cell_id"].apply(lambda x: " ".join(x)).reset_index()
sub_df.rename(columns={"cell_id": "cell_order"}, inplace=True)
sub_df.head()
sub_df.to_csv("submission.csv", index=False)
小结
- 本文的数据预处理方法是:将训练集的id、cell_id、cell_type、source、rank、ancestor_id、parent_id整合在一起作为属性列,然后为了方便预训练模型训练,将cell_type为markdown类型注释提取出来作为X,将rank归一化作为y。
- 在划分数据集时,使用到了GroupShuffleSplit,将祖先id作为组别信息,让模型在同组notebook中学习顺序关系。
- 通过在训练集上训练,让模型能够预测markdown在notebook中的位置。
reference
- https://www.kaggle.com/code/aerdem4/ai4code-pytorch-distilbert-baseline
- https://blog.csdn.net/y1040468929/article/details/124977901