notebooks/[UNUSED] Siamese Bert (Name + Attrs) Training.ipynb

{"cells":[{"cell_type":"code","execution_count":1,"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2024-08-21T08:44:24.847837Z","iopub.status.busy":"2024-08-21T08:44:24.847529Z","iopub.status.idle":"2024-08-21T08:44:32.116755Z","shell.execute_reply":"2024-08-21T08:44:32.115815Z","shell.execute_reply.started":"2024-08-21T08:44:24.847810Z"},"trusted":true},"outputs":[],"source":["import numpy as np\n","import pandas as pd \n","\n","import random\n","import json\n","import re\n","import gc\n","\n","import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","from torch.utils.data import DataLoader, Dataset\n","from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n","from transformers import AutoTokenizer, AutoModel, AdamW\n","from dataclasses import dataclass\n","\n","from sklearn.model_selection import train_test_split\n","\n","from tqdm import tqdm\n","tqdm.pandas()"]},{"cell_type":"code","execution_count":2,"metadata":{"execution":{"iopub.execute_input":"2024-08-21T08:44:32.120116Z","iopub.status.busy":"2024-08-21T08:44:32.119066Z","iopub.status.idle":"2024-08-21T08:44:32.126705Z","shell.execute_reply":"2024-08-21T08:44:32.125696Z","shell.execute_reply.started":"2024-08-21T08:44:32.120086Z"},"trusted":true},"outputs":[],"source":["@dataclass\n","class Config:\n","    model_name: str = 'cointegrated/rubert-tiny2'\n","    max_length: int = 768\n","    batch_size: int = 32   \n","    n_epochs: int = 1\n","    lr: float = 3e-5\n","    seed: int = 52\n","    test_size: float = 0.15\n","    \n","config = Config()"]},{"cell_type":"code","execution_count":3,"metadata":{"execution":{"iopub.execute_input":"2024-08-21T08:44:32.132689Z","iopub.status.busy":"2024-08-21T08:44:32.132082Z","iopub.status.idle":"2024-08-21T08:44:32.142414Z","shell.execute_reply":"2024-08-21T08:44:32.141396Z","shell.execute_reply.started":"2024-08-21T08:44:32.132654Z"},"trusted":true},"outputs":[],"source":["def seed_all(SEED):\n","    random.seed(SEED)\n","    np.random.seed(SEED)\n","    torch.manual_seed(SEED)\n","    torch.cuda.manual_seed(SEED)\n","    torch.backends.cudnn.deterministic = True\n","    \n","seed_all(config.seed)"]},{"cell_type":"code","execution_count":4,"metadata":{"execution":{"iopub.execute_input":"2024-08-21T08:44:32.144281Z","iopub.status.busy":"2024-08-21T08:44:32.143971Z","iopub.status.idle":"2024-08-21T08:47:33.525875Z","shell.execute_reply":"2024-08-21T08:47:33.524251Z","shell.execute_reply.started":"2024-08-21T08:44:32.144257Z"},"trusted":true},"outputs":[],"source":["text_and_bert = pd.read_parquet('/kaggle/input/extracted_data/text_and_bert.parquet', engine='pyarrow')\n","text_and_bert['description'] = text_and_bert['description'].fillna('no desc')"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["attrs = pd.read_parquet('/kaggle/input/extracted_data/attributes.parquet', columns=['characteristic_attributes_mapping', 'categories'], engine='pyarrow')\n","attrs['category_level_1'] = attrs['categories'].progress_apply(lambda x: eval(x)['1'])\n","attrs['category_level_2'] = attrs['categories'].progress_apply(lambda x: eval(x)['2'])"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["data = pd.concat([text_and_bert, attrs], axis=1)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["del text_and_bert, attrs\n","gc.collect()"]},{"cell_type":"code","execution_count":8,"metadata":{"execution":{"iopub.execute_input":"2024-08-21T08:47:46.068544Z","iopub.status.busy":"2024-08-21T08:47:46.067679Z","iopub.status.idle":"2024-08-21T08:49:21.933010Z","shell.execute_reply":"2024-08-21T08:49:21.931842Z","shell.execute_reply.started":"2024-08-21T08:47:46.068508Z"},"trusted":true},"outputs":[{"name":"stderr","output_type":"stream","text":["100%|██████████| 2252569/2252569 [01:35<00:00, 23523.02it/s]\n"]}],"source":["def create_attrs(row):\n","    if row['characteristic_attributes_mapping'] is not None:\n","        attributes = json.loads(row['characteristic_attributes_mapping'])\n","        attributes_str = '; '.join([f\"{key}: [{', '.join(values)}]\" for key, values in attributes.items()])\n","    else:\n","        attributes_str = 'none'\n","\n","    return attributes_str\n","\n","data['attrs'] = data.progress_apply(create_attrs, axis=1)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["data['text'] = data['name'].str.cat(data['category_level_1'], sep=' [SEP] ')\n","data['text'] = data['text'].str.cat(data['category_level_2'], sep=' [SEP] ')\n","data['text'] = data['text'].str.cat(data['attrs'], sep=' [SEP] ')"]},{"cell_type":"code","execution_count":5,"metadata":{"execution":{"iopub.execute_input":"2024-08-21T08:47:33.530132Z","iopub.status.busy":"2024-08-21T08:47:33.527415Z","iopub.status.idle":"2024-08-21T08:47:33.760610Z","shell.execute_reply":"2024-08-21T08:47:33.759617Z","shell.execute_reply.started":"2024-08-21T08:47:33.530084Z"},"trusted":true},"outputs":[],"source":["train_pairs = pd.read_parquet('/kaggle/input/extracted_data/train.parquet', engine='pyarrow')"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["train_data = data[['text', 'variantid']]"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.status.busy":"2024-08-21T08:47:35.120401Z","iopub.status.idle":"2024-08-21T08:47:35.120944Z","shell.execute_reply":"2024-08-21T08:47:35.120689Z","shell.execute_reply.started":"2024-08-21T08:47:35.120666Z"},"trusted":true},"outputs":[],"source":["train_pairs.rename(\n","    columns={\n","        'variantid1': 'variantid_1',\n","        'variantid2': 'variantid_2'\n","    }, inplace=True\n",")\n","\n","train_df = train_pairs.merge(\n","    train_data.add_suffix('_1'), \n","    on='variantid_1'\n",").merge(\n","    train_data.add_suffix('_2'), \n","    on='variantid_2'\n",")\n","\n","train_df = train_df[['text_1', 'text_2', 'target']]"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.status.busy":"2024-08-21T08:47:35.122580Z","iopub.status.idle":"2024-08-21T08:47:35.122963Z","shell.execute_reply":"2024-08-21T08:47:35.122798Z","shell.execute_reply.started":"2024-08-21T08:47:35.122782Z"},"trusted":true},"outputs":[],"source":["train_df, val_df = (\n","    train_test_split(\n","        train_df,\n","        test_size=config.test_size, \n","        random_state=config.seed, \n","        stratify=train_df.target)\n",")"]},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.status.busy":"2024-08-21T08:47:35.124320Z","iopub.status.idle":"2024-08-21T08:47:35.124703Z","shell.execute_reply":"2024-08-21T08:47:35.124552Z","shell.execute_reply.started":"2024-08-21T08:47:35.124538Z"},"trusted":true},"outputs":[],"source":["class TextPairDataset(Dataset):\n","    def __init__(self, df, tokenizer, max_length):\n","        self.df = df\n","        self.tokenizer = tokenizer\n","        self.max_length = max_length\n","\n","    def __len__(self):\n","        return len(self.df)\n","\n","    def __getitem__(self, idx):\n","        text_1 = self.df.iloc[idx]['text_1']\n","        text_2 = self.df.iloc[idx]['text_2']\n","        target = self.df.iloc[idx]['target']\n","\n","        encoding_1 = self.tokenizer.encode_plus(\n","            text_1[:1200], \n","            max_length=self.max_length, \n","            pad_to_max_length=False, \n","            return_attention_mask=True, \n","            return_tensors='pt', \n","            truncation=True\n","        )\n","\n","        encoding_2 = self.tokenizer.encode_plus(\n","            text_2[:1200], \n","            max_length=self.max_length, \n","            pad_to_max_length=False, \n","            return_attention_mask=True, \n","            return_tensors='pt', \n","            truncation=True\n","        )\n","\n","        return {\n","            'input_ids_1': encoding_1['input_ids'].flatten(),\n","            'attention_mask_1': encoding_1['attention_mask'].flatten(),\n","            'input_ids_2': encoding_2['input_ids'].flatten(),\n","            'attention_mask_2': encoding_2['attention_mask'].flatten(),\n","            'target': torch.tensor(target, dtype=torch.float)\n","        }"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["class SiameseBERT(nn.Module):\n","    def __init__(self):\n","        super(SiameseBERT, self).__init__()\n","        self.bert = AutoModel.from_pretrained('cointegrated/rubert-tiny2')\n","\n","    def forward(self, input_ids_1, attention_mask_1, input_ids_2, attention_mask_2):\n","        output_1 = self.bert(input_ids_1, attention_mask=attention_mask_1)\n","        output_2 = self.bert(input_ids_2, attention_mask=attention_mask_2)\n","        pooled_output_1 = output_1.pooler_output\n","        pooled_output_2 = output_2.pooler_output\n","        return pooled_output_1, pooled_output_2"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["tokenizer = AutoTokenizer.from_pretrained('cointegrated/rubert-tiny2')\n","\n","train_dataset = TextPairDataset(train_df, tokenizer, config.max_length)\n","val_dataset = TextPairDataset(val_df, tokenizer, config.max_length)\n","\n","train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)\n","val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["model = SiameseBERT()\n","criterion = nn.CosineEmbeddingLoss()\n","optimizer = optim.Adam(model.parameters(), lr=config.lr)\n","scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","model.to(device)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["def evaluate(model, val_loader, criterion, device):\n","    model.eval()\n","    val_loss = 0.0\n","    with torch.no_grad():\n","        for batch in val_loader:\n","            input_ids_1 = batch['input_ids_1'].to(device)\n","            attention_mask_1 = batch['attention_mask_1'].to(device)\n","            input_ids_2 = batch['input_ids_2'].to(device)\n","            attention_mask_2 = batch['attention_mask_2'].to(device)\n","            target = batch['target'].to(device)\n","\n","            output_1, output_2 = model(input_ids_1, attention_mask_1, input_ids_2, attention_mask_2)\n","            loss = criterion(output_1, output_2, target)\n","            val_loss += loss.item()\n","    return val_loss / len(val_loader)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["for epoch in range(config.n_epochs):\n","    model.train()\n","    running_loss = 0.0\n","    pbar = tqdm(train_loader, desc=f'Epoch {epoch + 1}', dynamic_ncols=True)\n","    for batch in pbar:\n","        input_ids_1 = batch['input_ids_1'].to(device)\n","        attention_mask_1 = batch['attention_mask_1'].to(device)\n","        input_ids_2 = batch['input_ids_2'].to(device)\n","        attention_mask_2 = batch['attention_mask_2'].to(device)\n","        target = batch['target'].to(device)\n","\n","        optimizer.zero_grad()\n","        output_1, output_2 = model(input_ids_1, attention_mask_1, input_ids_2, attention_mask_2)\n","        loss = criterion(output_1, output_2, target)\n","        loss.backward()\n","        optimizer.step()\n","        scheduler.step()\n","\n","        running_loss += loss.item()\n","        pbar.set_postfix({'Loss': loss.item()})\n","\n","    train_loss = running_loss / len(train_loader)\n","    val_loss = evaluate(model, val_loader, criterion, device)\n","    print(f'Epoch {epoch + 1}, Train Loss: {train_loss}, Val Loss: {val_loss}')\n","    \n","    torch.save(model.state_dict(), f'model_epoch_{epoch + 1}_valloss_{val_loss}.pth')"]}],"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"datasetId":3324098,"sourceId":5785473,"sourceType":"datasetVersion"},{"datasetId":5565627,"sourceId":9205012,"sourceType":"datasetVersion"},{"datasetId":5578618,"sourceId":9224239,"sourceType":"datasetVersion"}],"dockerImageVersionId":30747,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"}},"nbformat":4,"nbformat_minor":4}