{"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":"markdown","source":"**This notebook is a baseline for fine-tuning the `esm2` model.**\n\n- I used a `Colab Pro+ V100 high-memory instance` for training and inference.\n- I made some modifications to the evaluation code provided by the competition host (such as minor adjustments like float to float16).\n- The hyperparameter settings are arbitrary.\n\n\n**Reference code & discussions:**\n- https://www.kaggle.com/code/gusthema/cafa-5-protein-function-with-tensorflow\n- https://www.kaggle.com/code/henriupton/proteinet-pytorch-ems2-t5-protbert-embeddings\n- https://www.kaggle.com/code/sergeifironov/validate-ridge\n- https://www.kaggle.com/competitions/cafa-5-protein-function-prediction/discussion/407220\n- https://www.kaggle.com/code/yasufuminakama/fb3-deberta-v3-base-baseline-train (NLP pipeline)","metadata":{}},{"cell_type":"code","source":"# make submission (run only this cell)\nimport pandas as pd\nsubmission = pd.read_csv(\"/kaggle/input/cafa5-exp001/submission.tsv\", sep=\"\\t\", header=None)\nsubmission.to_csv(\"submission.tsv\",sep=\"\\t\", header=False, index=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-09T01:00:35.911822Z","iopub.execute_input":"2023-06-09T01:00:35.912952Z","iopub.status.idle":"2023-06-09T01:06:49.471671Z","shell.execute_reply.started":"2023-06-09T01:00:35.912906Z","shell.execute_reply":"2023-06-09T01:06:49.470516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Install","metadata":{}},{"cell_type":"code","source":"! pip install pyfastx -q\n! pip install polars\n! pip install transformers -q\n! pip install wandb -q\n! pip install 'git+https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup' -q\n! pip install torchmetrics -q","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nimport os\nimport json\nimport logging\nimport datetime\nimport requests\nimport gc\nimport joblib\nimport tempfile\nfrom pathlib import Path\n\nimport pyfastx\nimport pandas as pd\nimport numpy as np\nimport polars as pl\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nimport wandb\n\nfrom sklearn.model_selection import KFold\n\nfrom transformers import AutoTokenizer, AutoModel, AutoConfig\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.optim import Adam, AdamW\nfrom torch.optim.lr_scheduler import (\n    CosineAnnealingWarmRestarts,\n    CosineAnnealingLR,\n    MultiStepLR,\n    ReduceLROnPlateau,\n)\n\nfrom transformers import get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup\nfrom cosine_annealing_warmup import CosineAnnealingWarmupRestarts\nfrom transformers import DataCollatorWithPadding\nfrom torchmetrics.classification import MultilabelF1Score\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config","metadata":{}},{"cell_type":"code","source":"class Config:\n    # exp settings\n    name = \"exp001\"\n    competition = \"cafa-5-protein-function-prediction\"\n    debug = False\n    seed = 8823\n\n    train = True\n    evaluation = True\n    inference = True\n    submission = True\n\n    max_label = 512\n    n_fold = 5\n    trn_fold = [0]\n\n    max_epochs = 8\n    train_batch_size = 24\n    valid_batch_size = 64\n    num_workers = 4\n\n    model = \"facebook/esm2_t6_8M_UR50D\"\n    max_len = 512\n    gradient_checkpointing = False\n    gradient_accumulation_steps = 1\n    clip_grad_norm = 1000\n\n    optimizer = dict(\n        optimizer_name=\"AdamW\",\n        lr=2e-5,\n        weight_decay=1e-2,\n        eps=1e-6,\n        beta=(0.9, 0.999),\n        encoder_lr=2e-5,\n        decoder_lr=2e-5,\n    )\n\n    scheduler = dict(\n        scheduler_name=\"cosine_restarts\",\n        first_cycle_steps_ratio=0.5,\n        cycle_mult=1.0,\n        max_lr=2e-5,\n        min_lr=1e-7,\n        warmup_steps=100,\n        gamma=0.8,\n    )\n    batch_scheduler = True\n\n\nif Config.debug:\n    Config.max_label = 100\n    Config.max_epochs = 2\n    Config.max_len = 100\n    Config.n_fold = 2\n    Config.trn_fold = [0, 1]","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# constants\nHOME = Path(\"/content/drive/MyDrive/competition/kaggle_cafa-5-protein-function-prediction\")  # make appropriate modifications\nAPI = Path(\"/content/drive/MyDrive/competition/kaggle.json\")\n\nNOTEBOOK_NAME = requests.get(\"http://172.28.0.12:9000/api/sessions\").json()[0][\"name\"][:-6]\nEXP_NAME = Config.name if Config.name is not None else NOTEBOOK_NAME\nINPUTS = HOME / \"inputs\"\nOUTPUTS = HOME / \"outputs\"\nINTERMIDIATES = HOME / \"intermidiates\"  # intermidiate outputs\nSUBMISSIONS = HOME / \"submissions\"\nOUTPUTS_EXP = OUTPUTS / EXP_NAME\nEXP_MODELS = OUTPUTS_EXP / \"models\"\nEXP_REPORTS = OUTPUTS_EXP / \"reports\"\nEXP_PREDS = OUTPUTS_EXP / \"predictions\"\n\nSCRIPTS = HOME / \"scripts\"\nCAFA_EVALUATOR = SCRIPTS / \"CAFA_evaluator\"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Setup","metadata":{}},{"cell_type":"code","source":"def setup(Config):\n    print(\"This environment is Google Colab\")\n\n    # mount\n    from google.colab import drive\n\n    CONTENT = Path(\"/content/drive\")\n    if not CONTENT.is_dir():\n        drive.mount(CONTENT.as_posix())\n\n    # use kaggle api (need kaggle token)\n    with open(API, \"r\") as f:\n        json_data = json.load(f)\n    os.environ[\"KAGGLE_USERNAME\"] = json_data[\"username\"]\n    os.environ[\"KAGGLE_KEY\"] = json_data[\"key\"]\n\n    for d in [\n        HOME,\n        INPUTS,\n        SUBMISSIONS,\n        EXP_MODELS,\n        EXP_REPORTS,\n        EXP_PREDS,\n        INTERMIDIATES,\n        SCRIPTS,\n        CAFA_EVALUATOR,\n    ]:\n        d.mkdir(parents=True, exist_ok=True)\n\n    # download datasets\n    if not check_file_exists(folder_path=INPUTS, file_name=\"sample_submission\"):\n        print(\"start downloading datasets 🚀\")\n        ! pip install --upgrade --force-reinstall --no-deps kaggle\n        ! kaggle competitions download -c $Config.competition -p $INPUTS\n        filepath = INPUTS / f\"{Config.competition}.zip\"\n        ! unzip -d $INPUTS $filepath\n    \n    # clone evaluator code\n    if not (CAFA_EVALUATOR / \"src\").is_dir():\n        ! git clone https://github.com/BioComputingUP/CAFA-evaluator.git $CAFA_EVALUATOR\n\n\ndef check_file_exists(folder_path, file_name):\n    folder = Path(folder_path)\n    for file_path in folder.glob(\"**/*\"):\n        if file_path.is_file() and file_path.stem == file_name:\n            return True\n    return False\n\n\ndef create_new_datasets_in_kaggle(dataset_name, upload_dir):\n    from kaggle.api.kaggle_api_extended import KaggleApi\n\n    dataset_metadata = {}\n    dataset_metadata[\"id\"] = f'{os.environ[\"KAGGLE_USERNAME\"]}/{dataset_name}'\n    dataset_metadata[\"licenses\"] = [{\"name\": \"CC0-1.0\"}]\n    dataset_metadata[\"title\"] = dataset_name\n    with open(os.path.join(upload_dir, \"dataset-metadata.json\"), \"w\") as f:\n        json.dump(dataset_metadata, f, indent=4)\n    api = KaggleApi()\n    api.authenticate()\n    api.dataset_create_new(folder=upload_dir, convert_to_csv=False, dir_mode=\"tar\")\n\n    \ndef upload_submission_file_to_kaggle(dataset_name, submission_df):\n    \"\"\"\n    upload a CAFA5 submission file to kaggle dataset (use temporary directory)\n    \"\"\"\n    with tempfile.TemporaryDirectory() as tmp_dir:\n        dirpath = Path(tmp_dir)\n        submission_df.to_csv(dirpath / \"submission.tsv\", header=False, index=False, sep=\"\\t\")\n        create_new_datasets_in_kaggle(dataset_name=dataset_name, upload_dir=dirpath)\n\n        \nclass Logger:\n    def __init__(self, path):\n        self.general_logger = logging.getLogger(path)\n        stream_handler = logging.StreamHandler()\n        file_general_handler = logging.FileHandler(os.path.join(path, \"Experiment.log\"))\n        if len(self.general_logger.handlers) == 0:\n            self.general_logger.addHandler(stream_handler)\n            self.general_logger.addHandler(file_general_handler)\n            self.general_logger.setLevel(logging.INFO)\n\n    def info(self, message):\n        # display time\n        self.general_logger.info(\"[{}] - {}\".format(self.now_string(), message))\n\n    @staticmethod\n    def now_string():\n        return str(datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\"))\n\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"setup(Config)\nLOGGER = Logger(OUTPUTS_EXP.as_posix())\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nwandb.login()  # need wandb account\nsys.path.append((CAFA_EVALUATOR / \"src\").as_posix())  # use evaluator as package","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Metrics","metadata":{}},{"cell_type":"code","source":"# competition metrics\ndef get_score(gt_df, pred_df, th_step=0.05, max_terms=None, use_cache=False, n_cpu=1):\n    from parser import obo_parser, gt_parser, pred_parser, ia_parser\n    from graph import Graph\n    from evaluation import evaluate_prediction\n\n    pred_df = pred_df.astype({\"preds\": np.float16})\n    gt_df = gt_df[[\"EntryID\", \"term\"]]\n\n    with tempfile.TemporaryDirectory() as tmp_dir:\n        dirpath = Path(tmp_dir)\n\n        # set filepath and parameter\n        gt_file, pred_file = dirpath / \"gt.tsv\", dirpath / \"pred.tsv\"\n        ia = INPUTS / \"IA.txt\"\n        obo_file = INPUTS / \"Train\" / \"go-basic.obo\"\n        no_orphans = False\n        prop = \"max\"\n        norm = \"cafa\"\n        parsed_gt_file = INTERMIDIATES / \"parsed_gt.pkl\"\n\n        # save evaluator inputs to tmpdir\n        gt_df.to_csv(gt_file, header=False, index=False, sep=\"\\t\")\n        pred_df.to_csv(pred_file, header=False, index=False, sep=\"\\t\")\n\n        # main\n        ia_dict = None if ia is None else ia_parser(ia)\n        tau_arr = np.arange(0.01, 1, th_step)\n        ontologies = [\n            Graph(ns, terms_dict, ia_dict, not no_orphans)\n            for ns, terms_dict in obo_parser(obo_file).items()\n        ]\n\n        if use_cache:\n            if not parsed_gt_file.is_file():\n                # to save time\n                gt = gt_parser(gt_file, ontologies)\n                joblib.dump(gt, parsed_gt_file)\n            gt = joblib.load(parsed_gt_file)\n        else:\n            gt = gt_parser(gt_file, ontologies)\n\n        df_pred = evaluate_prediction(  # too heavy and slow...\n            prediction=pred_parser(pred_file, ontologies, gt, prop_mode=prop, max_terms=max_terms),\n            gt=gt,\n            ontologies=ontologies,\n            tau_arr=tau_arr,\n            normalization=norm,\n            n_cpu=n_cpu,\n        )\n    score = df_pred.groupby(\"ns\").agg({\"f\": \"max\"}).values.mean()\n    return score, df_pred.reset_index(drop=True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Data","metadata":{}},{"cell_type":"code","source":"def make_train_terms_crosstab(train_terms_df):\n    filepath = INTERMIDIATES / \"train_terms_crosstab.pkl\"\n\n    if not filepath.is_file():\n        crosstab_df = train_terms_df.pivot(\n            values=\"aspect\", index=\"EntryID\", columns=\"term\", aggregate_function=\"count\"\n        ).fill_null(0)\n        joblib.dump(\n            crosstab_df.to_pandas().astype(bool).assign(EntryID=crosstab_df[\"EntryID\"]),\n            filepath,\n        )\n\n        del crosstab_df\n        gc.collect()\n\n    # use polars\n    return joblib.load(filepath).sort_values(\"EntryID\").reset_index(drop=True)\n\n\ndef make_sequence_df(input_fasta):\n    sequence_df = (\n        pd.DataFrame([[x.name, x.seq] for x in input_fasta], columns=[\"EntryID\", \"seq\"])\n        .sort_values(\"EntryID\")\n        .reset_index(drop=True)\n    )\n    return sequence_df\n\n\ndef make_train_dataset(train_seq_df, train_label_df, max_label=None):\n    train_label_df = train_label_df.set_index(\"EntryID\")\n    if max_label is None:\n        max_label = train_label_df.shape[1]\n\n    filepath = INTERMIDIATES / \"sorted_GO_list.pkl\"\n    if not filepath.is_file():\n        target_columns = train_label_df.sum(axis=0).sort_values(ascending=False).index.tolist()\n        joblib.dump(target_columns, filepath)\n\n    target_columns = joblib.load(filepath)[:max_label]\n    train_df = pd.concat(\n        [train_seq_df.set_index(\"EntryID\"), train_label_df[target_columns]], axis=1\n    )\n    return train_df, target_columns","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# raw data\ntrain_terms_df = pl.read_csv(INPUTS / \"Train/train_terms.tsv\", separator=\"\\t\")\nsample_sumbmission_df = pl.read_csv(\n    INPUTS / \"sample_submission.tsv\", separator=\"\\t\"\n)  # [\"Protein Id\", \"GO Term Id\", \"Prediction\"]\n\ntrain_fasta = pyfastx.Fasta((INPUTS / \"Train\" / \"train_sequences.fasta\").as_posix())\ntest_fasta = pyfastx.Fasta((INPUTS / \"Test (Targets)\" / \"testsuperset.fasta\").as_posix())\n\n# preprocess\ntrain_label_df = make_train_terms_crosstab(train_terms_df)  # EntryID * all GO:\ntrain_seq_df = make_sequence_df(input_fasta=train_fasta)\ntest_df = make_sequence_df(input_fasta=test_fasta)\ntrain_df, target_columns = make_train_dataset(\n    train_seq_df, train_label_df, max_label=Config.max_label\n)\ntrain_terms_df = train_terms_df.to_pandas()\ntest_df = test_df.set_index(\"EntryID\")\n\ndel train_fasta, test_fasta, train_seq_df, train_label_df\ngc.collect()\n\n# load tokenizer\ntokenizer = AutoTokenizer.from_pretrained(Config.model)\n\nif Config.debug:\n    train_df = train_df.sample(1000, random_state=Config.seed)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CV Split","metadata":{}},{"cell_type":"code","source":"def add_fold_idx(config, train_df, target_columns):\n    fold = KFold(n_splits=config.n_fold, shuffle=True, random_state=config.seed)\n    train_df[\"fold\"] = -1\n    for i_fold, (train_index, val_index) in enumerate(\n        fold.split(train_df, train_df[target_columns])\n    ):\n        train_df.iloc[val_index, train_df.columns.get_loc(\"fold\")] = int(i_fold)\n    train_df[\"fold\"] = train_df[\"fold\"].astype(int)\n    return train_df","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = add_fold_idx(config=Config, train_df=train_df, target_columns=target_columns)\ntrain_df.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"code","source":"def find_max_length(tokenizer, texts: list[str]):\n    \"\"\"\n    [memo]\n    train : facebook/esm2_t6_8M_UR50D : 35377\n    \"\"\"\n    lengths = [\n        len(tokenizer(text, add_special_tokens=False)[\"input_ids\"])\n        for text in tqdm(texts, total=len(texts))\n    ]\n    return lengths\n\n\ndef prepare_input(config, tokenizer, text):\n    inputs = tokenizer.encode_plus(\n        text,\n        return_tensors=None,\n        add_special_tokens=True,\n        max_length=config.max_len,\n        padding=\"max_length\",\n        truncation=True,\n    )\n    for k, v in inputs.items():\n        inputs[k] = torch.tensor(v, dtype=torch.long)\n    return inputs\n\n\nclass TrainDataset(Dataset):\n    def __init__(self, config, tokenizer, texts: list[str], targets: list | np.ndarray):\n        self.config = config\n        self.texts = texts\n        self.targets = targets\n        self.tokenizer = tokenizer\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, item):\n        inputs = prepare_input(config=self.config, tokenizer=self.tokenizer, text=self.texts[item])\n        targets = torch.tensor(self.targets[item], dtype=torch.float16)  # label\n        return {\n            \"input_ids\": inputs[\"input_ids\"],\n            \"attention_mask\": inputs[\"attention_mask\"],\n            \"targets\": targets,\n        }\n\n\nclass TestDataset(Dataset):\n    def __init__(self, config, tokenizer, texts: list[str]):\n        self.config = config\n        self.texts = texts\n        self.tokenizer = tokenizer\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, item):\n        inputs = prepare_input(config=self.config, tokenizer=self.tokenizer, text=self.texts[item])\n        return inputs\n\n\ndef collate(inputs):\n    mask_len = int(inputs[\"attention_mask\"].sum(axis=1).max())\n    for k, v in inputs.items():\n        if k != \"targets\":\n            inputs[k] = inputs[k][:, :mask_len]\n    return inputs","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"def get_model_config(config):\n    model_config = AutoConfig.from_pretrained(config.model, output_hidden_states=True)\n    model_config.hidden_dropout = 0.0\n    model_config.hidden_dropout_prob = 0.0\n    model_config.attention_dropout = 0.0\n    model_config.attention_probs_dropout_prob = 0.0\n    return model_config\n\n\nclass CustomModel(nn.Module):\n    def __init__(self, config, num_label, pretrained=False):\n        super().__init__()\n        self.model_config = get_model_config(config=config)\n\n        if pretrained:\n            self.model = AutoModel.from_pretrained(config.model, config=self.model_config)\n        else:\n            self.model = AutoModel.from_config(self.model_config)\n\n        if config.gradient_checkpointing:\n            self.model.gradient_checkpointing_enable()\n\n        self.fc = nn.Linear(self.model_config.hidden_size, num_label)\n        self._init_weights(self.fc)\n\n    def _init_weights(self, module):\n        if isinstance(module, nn.Linear):\n            module.weight.data.normal_(mean=0.0, std=self.model_config.initializer_range)\n            if module.bias is not None:\n                module.bias.data.zero_()\n        elif isinstance(module, nn.Embedding):\n            module.weight.data.normal_(mean=0.0, std=self.model_config.initializer_range)\n            if module.padding_idx is not None:\n                module.weight.data[module.padding_idx].zero_()\n        elif isinstance(module, nn.LayerNorm):\n            module.bias.data.zero_()\n            module.weight.data.fill_(1.0)\n\n    def feature(self, input_ids, attention_mask):\n        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)\n        feature = outputs[0][:, 0, :]\n        return feature\n\n    def forward(self, batch):\n        feature = self.feature(batch[\"input_ids\"], batch[\"attention_mask\"])\n        output = self.fc(feature)\n        return output","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Funcs","metadata":{}},{"cell_type":"code","source":"def train_fn(\n    config,\n    model,\n    dataloader,\n    criterion,\n    optimizer,\n    scheduler,\n    device,\n    wandb_logger,\n    _custom_step,\n):\n    model.train()\n    scaler = torch.cuda.amp.GradScaler()\n    losses = []\n\n    tbar = tqdm(enumerate(dataloader), total=len(dataloader))\n    for step, batch in tbar:\n        batch = collate(batch)\n        for k, v in batch.items():\n            batch[k] = v.to(device)\n        targets = batch[\"targets\"]\n        batch_size = targets.size(0)\n\n        with torch.cuda.amp.autocast():\n            batch_outputs = model(batch)\n            loss = criterion(batch_outputs, targets)\n\n        if config.gradient_accumulation_steps > 1:\n            loss = loss / config.gradient_accumulation_steps\n\n        scaler.scale(loss).backward()\n        if config.clip_grad_norm is not None:\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config.clip_grad_norm)\n\n        if (step + 1) % config.gradient_accumulation_steps == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            _custom_step += 1\n            if config.batch_scheduler:\n                scheduler.step()\n\n        wandb_logger.log(\n            {\"train_loss\": loss, \"lr\": scheduler.get_lr()[0], \"train_step\": _custom_step}\n        )\n\n        losses.append(float(loss))\n        tbar.set_description(f\"loss: {np.mean(losses):.4f} lr: {scheduler.get_lr()[0]:.6f}\")\n\n    loss = np.mean(losses)\n    return loss, _custom_step\n\n\ndef valid_fn(\n    config,\n    model,\n    dataloader,\n    criterion,\n    device,\n    wandb_logger,\n    _custom_step,\n):\n    model.eval()\n    outputs, targets = [], []\n    losses = []\n\n    tbar = tqdm(enumerate(dataloader), total=len(dataloader))\n    for step, batch in tbar:\n        batch = collate(batch)\n        targets.append(batch[\"targets\"])\n\n        for k, v in batch.items():\n            batch[k] = v.to(device)\n\n        batch_size = batch[\"targets\"].size(0)\n        with torch.no_grad():\n            batch_outputs = model(batch)\n            loss = criterion(batch_outputs, batch[\"targets\"])\n\n        if config.gradient_accumulation_steps > 1:\n            loss = loss / config.gradient_accumulation_steps\n\n        batch_outputs = torch.sigmoid(batch_outputs)  # to proba\n        batch_outputs = batch_outputs.to(\"cpu\").numpy()\n        outputs.append(batch_outputs)\n\n        wandb_logger.log({\"valid_loss\": loss, \"valid_step\": _custom_step})\n        _custom_step += 1\n        losses.append(float(loss))\n\n        tbar.set_description(f\"loss: {np.mean(losses):.4f}\")\n\n    outputs = np.concatenate(outputs)\n    targets = np.concatenate(targets)\n\n    loss = np.mean(losses)\n    return (loss, outputs, targets, _custom_step)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_optimizer(optimizer_config, model):\n    if optimizer_config[\"optimizer_name\"] == \"AdamW\":\n        param_optimizer = list(model.named_parameters())\n        no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n        optimizer_parameters = [\n            {\n                \"params\": [\n                    p\n                    for n, p in model.model.named_parameters()\n                    if not any(nd in n for nd in no_decay)\n                ],\n                \"lr\": optimizer_config[\"encoder_lr\"],\n                \"weight_decay\": optimizer_config[\"weight_decay\"],\n            },\n            {\n                \"params\": [\n                    p for n, p in model.model.named_parameters() if any(nd in n for nd in no_decay)\n                ],\n                \"lr\": optimizer_config[\"encoder_lr\"],\n                \"weight_decay\": 0.0,\n            },\n            {\n                \"params\": [p for n, p in model.named_parameters() if \"model\" not in n],\n                \"lr\": optimizer_config[\"decoder_lr\"],\n                \"weight_decay\": 0.0,\n            },\n        ]\n        optimizer = AdamW(\n            optimizer_parameters,\n            lr=optimizer_config[\"lr\"],\n            betas=optimizer_config[\"beta\"],\n            eps=optimizer_config[\"eps\"],\n        )\n        return optimizer\n    else:\n        raise NotImplementedError\n\n\ndef get_scheduler(scheduler_config, optimizer, num_train_steps):\n    if scheduler_config[\"scheduler_name\"] == \"linear\":\n        scheduler = get_linear_schedule_with_warmup(\n            optimizer,\n            num_warmup_steps=int(scheduler_config[\"num_warmup_steps_rate\"] * num_train_steps),\n            num_training_steps=num_train_steps,\n        )\n        return scheduler\n\n    elif scheduler_config[\"scheduler_name\"] == \"cosine\":\n        scheduler = get_cosine_schedule_with_warmup(\n            optimizer,\n            num_warmup_steps=int(scheduler_config[\"num_warmup_steps_rate\"] * num_train_steps),\n            num_training_steps=num_train_steps,\n            num_cycles=scheduler_config[\"num_cycles\"],\n        )\n        return scheduler\n\n    elif scheduler_config[\"scheduler_name\"] == \"cosine_restarts\":\n        \"\"\"\n        example:\n            first_cycle_steps_ratio = 0.25,\n            cycle_mult = 1.0,\n            max_lr = 2e-5,\n            min_lr = 1e-7,\n            warmup_steps=100,\n            gamma=0.8)\n        \"\"\"\n        scheduler = CosineAnnealingWarmupRestarts(\n            optimizer,\n            first_cycle_steps=int(num_train_steps * scheduler_config[\"first_cycle_steps_ratio\"]),\n            cycle_mult=scheduler_config[\"cycle_mult\"],\n            max_lr=scheduler_config[\"max_lr\"],\n            min_lr=scheduler_config[\"min_lr\"],\n            warmup_steps=scheduler_config[\"warmup_steps\"],\n            gamma=scheduler_config[\"gamma\"],\n        )\n        return scheduler\n\n    else:\n        raise NotImplementedError\n\n\ndef train_loop(config, name, train_df, valid_df, target_columns, tokenizer, device):\n    LOGGER.info(f\"========== {name} training ==========\")\n\n    # set wandb logger\n    wandb.init(\n        project=config.competition,\n        name=name,\n        group=f\"{config.name}\",\n        job_type=\"train\",\n        anonymous=None,\n        reinit=True,\n    )\n\n    # dataset, dataloader\n    train_dataset = TrainDataset(\n        config=config,\n        texts=train_df[\"seq\"].tolist(),\n        targets=train_df[target_columns].to_numpy(),\n        tokenizer=tokenizer,\n    )\n    valid_dataset = TrainDataset(\n        config=config,\n        texts=valid_df[\"seq\"].tolist(),\n        targets=valid_df[target_columns].to_numpy(),\n        tokenizer=tokenizer,\n    )\n    train_dataloader = DataLoader(\n        train_dataset,\n        batch_size=config.train_batch_size,\n        shuffle=True,\n        num_workers=config.num_workers,\n        pin_memory=True,\n        drop_last=True,\n    )\n    valid_dataloader = DataLoader(\n        valid_dataset,\n        batch_size=config.valid_batch_size,\n        shuffle=False,\n        num_workers=config.num_workers,\n        pin_memory=True,\n        drop_last=False,\n    )\n\n    # set model & optimizer\n    model = CustomModel(config=config, num_label=len(target_columns), pretrained=True)\n    model.to(device)\n    optimizer = get_optimizer(optimizer_config=config.optimizer, model=model)\n\n    # set scheduler\n    num_train_steps = int(\n        len(train_dataloader) * config.max_epochs // config.gradient_accumulation_steps\n    )\n    scheduler = get_scheduler(\n        scheduler_config=config.scheduler, optimizer=optimizer, num_train_steps=num_train_steps\n    )\n\n    # loop\n    criterion = nn.BCEWithLogitsLoss()\n    score_fn = MultilabelF1Score(num_labels=len(target_columns))  # TODO : use custom f1\n\n    best_score = -np.inf\n    tr_step, val_step = 0, 0\n    for epoch in range(Config.max_epochs):\n        # training\n        loss, tr_step = train_fn(\n            config=config,\n            model=model,\n            dataloader=train_dataloader,\n            criterion=criterion,\n            optimizer=optimizer,\n            scheduler=scheduler,\n            device=device,\n            wandb_logger=wandb,\n            _custom_step=tr_step,\n        )\n\n        # validation\n        val_loss, val_outputs, val_targets, val_step = valid_fn(\n            config=config,\n            model=model,\n            dataloader=valid_dataloader,\n            criterion=criterion,\n            device=device,\n            wandb_logger=wandb,\n            _custom_step=val_step,\n        )\n\n        # calc score (the higher, the better)\n        eval_score = score_fn(torch.tensor(val_outputs).float(), torch.tensor(val_targets).float())\n        logs = {\n            \"Epoch\": epoch,\n            \"eval_score\": eval_score,\n            \"train_loss_epoch\": loss.item(),\n            \"valid_loss_epoch\": val_loss.item(),\n        }\n        wandb.log(logs)\n        LOGGER.info(logs)\n\n        if best_score < eval_score:\n            best_score = eval_score\n            LOGGER.info(f\"epoch {epoch} - best score: {best_score:.4f} model\")\n\n            torch.save(model.state_dict(), EXP_MODELS / f\"{name}.pth\")  # save model weight\n            joblib.dump(val_outputs, EXP_PREDS / f\"{name}.pkl\")  # save outputs\n\n        if not config.batch_scheduler:\n            scheduler.step()\n\n    torch.cuda.empty_cache()\n    gc.collect()\n    wandb.finish(quiet=True)\n\n    # save best predictions with id\n    best_val_outputs = joblib.load(EXP_PREDS / f\"{name}.pkl\")\n    outputs = {\n        \"id\": valid_df.index.tolist(),\n        \"target_names\": target_columns,\n        \"predictions\": best_val_outputs,\n    }\n    joblib.dump(outputs, EXP_PREDS / f\"{name}_best.pkl\")\n\n    return outputs","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def inference_fn(test_dataloader, model, device):\n    preds, targets_masks = [], []\n    model.eval()\n    model.to(device)\n\n    tbar = tqdm(test_dataloader, total=len(test_dataloader))\n    for batch in tbar:\n        for k, v in batch.items():\n            batch[k] = v.to(device)\n\n        with torch.no_grad():\n            outputs = model(batch)\n            outputs = torch.sigmoid(outputs)\n\n        outputs = outputs.cpu().detach().numpy()\n        preds.append(outputs)\n\n    return np.concatenate(preds)\n\n\ndef get_predictions(config, test_df, tokenizer, model_path, target_columns, device):\n    test_dataset = TestDataset(config=config, tokenizer=tokenizer, texts=test_df[\"seq\"])\n    test_dataloader = DataLoader(\n        test_dataset,\n        batch_size=config.valid_batch_size,\n        shuffle=False,\n        collate_fn=DataCollatorWithPadding(tokenizer=tokenizer, padding=\"longest\"),\n        num_workers=config.num_workers,\n        pin_memory=True,\n        drop_last=False,\n    )\n\n    # get model\n    model = CustomModel(config=config, num_label=len(target_columns), pretrained=False)\n    state = torch.load(model_path)\n    model.load_state_dict(state)\n    predictions = inference_fn(test_dataloader, model, device)\n\n    del model, state, test_dataloader\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    outputs = {\n        \"id\": test_df.index.tolist(),\n        \"target_names\": target_columns,\n        \"predictions\": predictions,\n    }\n    return outputs","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"if Config.train:\n\n    for i_fold in range(Config.n_fold):\n        if i_fold not in Config.trn_fold:\n            continue\n\n        train_df_fold = train_df[train_df[\"fold\"] != i_fold]\n        valid_df_fold = train_df[train_df[\"fold\"] == i_fold]\n\n        train_loop(\n            config=Config,\n            name=f\"fold_{i_fold}\",\n            train_df=train_df_fold,\n            valid_df=valid_df_fold,\n            target_columns=target_columns,\n            tokenizer=tokenizer,\n            device=DEVICE,\n        )","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Evaluation","metadata":{}},{"cell_type":"code","source":"def make_submission_formatted_outputs_df(outputs: dict) -> pd.DataFrame:\n    data = {\n        \"id\": np.repeat(outputs[\"id\"], len(outputs[\"target_names\"])),\n        \"terms\": np.tile(outputs[\"target_names\"], len(outputs[\"id\"])),\n        \"preds\": outputs[\"predictions\"].ravel(),\n    }\n    submission_df = pd.DataFrame(data)\n    return submission_df\n\n\ndef make_crosstab_formatted_outputs_df(outputs: dict) -> pd.DataFrame:\n    output_df = pd.DataFrame(\n        outputs[\"predictions\"],\n        columns=outputs[\"target_names\"],\n        index=pd.Index(outputs[\"id\"], name=\"id\"),\n    )\n    return output_df\n\nif Config.evaluation:\n    pred_dfs, gt_dfs = [], []\n    for i_fold in range(Config.n_fold):\n        if i_fold not in Config.trn_fold:\n            continue\n\n        # get data\n        outputs = joblib.load(EXP_PREDS / f\"fold_{i_fold}_best.pkl\")\n        pred_df = make_submission_formatted_outputs_df(outputs)\n        gt_df = train_terms_df[train_terms_df[\"EntryID\"].isin(pred_df[\"id\"])].reset_index(drop=True)\n\n        # calc score\n        score, score_df = get_score(gt_df=gt_df, pred_df=pred_df, th_step=0.1)\n\n        # save result\n        LOGGER.info(f\"fold {i_fold} : score = {score:.4f}\")\n        score_df.to_csv(EXP_REPORTS / f\"score_df_fold_{i_fold}.csv\", index=False)\n\n        # stack\n        pred_dfs.append(pred_df)\n        gt_dfs.append(gt_df)\n\n    all_pred_df = pd.concat(pred_dfs, axis=0).reset_index(drop=True)\n    all_gt_df = pd.concat(gt_dfs, axis=0).reset_index(drop=True)\n    all_score, all_score_df = get_score(gt_df=all_gt_df, pred_df=all_pred_df, th_step=0.1)\n    LOGGER.info(f\"all : score = {all_score:.4f}\")\n\n    joblib.dump(all_pred_df, EXP_PREDS / \"all_pred_df.pkl\")\n    joblib.dump(all_gt_df, EXP_PREDS / \"all_gt_df.pkl\")\n    joblib.dump(all_score_df, EXP_PREDS / \"all_score_df.pkl\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"code","source":"if Config.inference:\n    test_pred_df = pd.DataFrame()\n    predictions = 0\n    n_model = 0\n    for i_fold in range(Config.n_fold):\n        if i_fold not in Config.trn_fold:\n            continue\n\n        outputs = get_predictions(\n            config=Config, \n            test_df=test_df, \n            tokenizer=tokenizer, \n            model_path=EXP_MODELS / f\"fold_{i_fold}.pth\", \n            target_columns=target_columns, \n            device=DEVICE)\n        predictions += outputs[\"predictions\"]\n        n_model += 1\n\n    outputs[\"predictions\"] = predictions / n_model  # fold ave, same id & targets\n    test_pred_df = make_submission_formatted_outputs_df(outputs)\n    test_pred_df.reset_index(drop=True, inplace=True)\n    joblib.dump(test_pred_df, EXP_PREDS / \"test_pred_df.pkl\")\n    test_pred_df.to_csv(SUBMISSIONS / f\"{Config.name}.tsv\", header=False, index=False, sep=\"\\t\")\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submission","metadata":{}},{"cell_type":"code","source":"if Config.submission:\n    submission_df = joblib.load(EXP_PREDS / \"test_pred_df.pkl\")\n    upload_submission_file_to_kaggle(dataset_name=f\"CAFA5-{Config.name}\", submission_df=submission_df)","metadata":{},"execution_count":null,"outputs":[]}]}