{"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":"## preprocess","metadata":{}},{"cell_type":"code","source":"import json\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nfrom scipy import sparse\nfrom tqdm import tqdm\n\nfrom transformers import AutoTokenizer\n\nfrom typing import Optional, Tuple, Union\n\nimport torch\nfrom torch import nn\nfrom torch.nn import CrossEntropyLoss\nfrom transformers.modeling_outputs import SequenceClassifierOutput\nfrom transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaModel\n\nimport sys, os\nfrom transformers import AutoModel\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.nn.functional as F\n\n\npd.options.display.width = 180\npd.options.display.max_colwidth = 120\ndata_dir = Path('../input/AI4Code')\n\ndef read_notebook(path):\n    return (\n        pd.read_json(\n            path,\n            dtype={'cell_type': 'category', 'source': 'str'})\n        .assign(id=path.stem)\n        .rename_axis('cell_id')\n    )\n\npaths_test = list((data_dir / 'test').glob('*.json'))\n# paths_test = list((data_dir / 'train').glob('*.json'))[:20000]\nnotebooks_test = [\n    read_notebook(path) for path in tqdm(paths_test, desc='Test NBs')\n]\ntest_df = (\n    pd.concat(notebooks_test)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n).reset_index()\ntest_df[\"rank\"] = test_df.groupby([\"id\", \"cell_type\"]).cumcount()\ntest_df[\"pred\"] = test_df.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=True)\n\n# Additional code cells\ndef clean_code(cell):\n    return str(cell).replace(\"\\\\n\", \"\\n\")\n\n\ndef sample_cells(cells, n, ext):\n    cells = [clean_code(cell) for cell in cells]\n    if n >= len(cells):\n        return [cell[:200] for cell in cells]\n    else:\n        results = []\n        step = len(cells) / n\n        idx = 0\n        while int(np.round(idx)) < len(cells):\n            results.append(cells[int(np.round(idx))])\n            if idx == 0:\n                idx += ext\n            idx += step\n        assert cells[0] in results\n        if cells[-1] not in results:\n            results[-1] = cells[-1]\n        return results\n\n\ndef get_features(df, ext):\n    features = dict()\n    df = df.sort_values(\"rank\").reset_index(drop=True)\n    for idx, sub_df in tqdm(df.groupby(\"id\")):\n        features[idx] = dict()\n        total_md = sub_df[sub_df.cell_type == \"markdown\"].shape[0]\n        code_sub_df = sub_df[sub_df.cell_type == \"code\"]\n        total_code = code_sub_df.shape[0]\n        codes = sample_cells(code_sub_df.source.values, 30, ext)\n        features[idx][\"total_code\"] = total_code\n        features[idx][\"total_md\"] = total_md\n        features[idx][\"codes\"] = codes\n    return features\n\ntest_fts1 = get_features(test_df, 0)\ntest_fts2 = get_features(test_df, 1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 一段階目 ","metadata":{}},{"cell_type":"markdown","source":"pairwise","metadata":{}},{"cell_type":"code","source":"model_paths = [\n    \"../input/ai4code-pairwise-model/codebert-1epochKTKA-4epochKA\",\n]\n\nmax_length = 128\ntemperature = 0.7\n\ndf_test = (\n    pd.concat(notebooks_test)\n    .set_index('id', append=True)\n    .swaplevel()\n    .sort_index(level='id', sort_remaining=False)\n)\n\ndf_test = df_test.reset_index()\n\n\nclass MarkdownModelPairwise(RobertaPreTrainedModel):\n    def __init__(self, config):\n        super().__init__(config)\n\n        self.roberta = RobertaModel(config, add_pooling_layer=False)\n\n        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n        self.activation = nn.Tanh()\n\n        # Initialize weights and apply final processing\n        self.post_init()\n\n    def forward(\n        self,\n        input_ids: Optional[torch.LongTensor] = None,\n        attention_mask: Optional[torch.FloatTensor] = None,\n        token_type_ids: Optional[torch.LongTensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n        head_mask: Optional[torch.FloatTensor] = None,\n        inputs_embeds: Optional[torch.FloatTensor] = None,\n        labels: Optional[torch.LongTensor] = None,\n        output_attentions: Optional[bool] = None,\n        output_hidden_states: Optional[bool] = None,\n        return_dict: Optional[bool] = None,\n\n        input_ids_code: Optional[torch.LongTensor] = None,\n        attention_mask_code: Optional[torch.FloatTensor] = None,\n    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:\n        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n\n        outputs_markdown = self.roberta(\n            input_ids,\n            attention_mask=attention_mask,\n            token_type_ids=token_type_ids,\n            position_ids=position_ids,\n            head_mask=head_mask,\n            inputs_embeds=inputs_embeds,\n            output_attentions=output_attentions,\n            output_hidden_states=output_hidden_states,\n            return_dict=return_dict,\n        )\n        sequence_output_markdown = outputs_markdown[0]\n        # logits = self.classifier(sequence_output)\n\n        # pooled_output = sequence_output[:, 0, :]\n        pooled_output_markdown = (sequence_output_markdown * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.unsqueeze(-1).sum(1)\n        pooled_output_markdown = self.dense(pooled_output_markdown)\n        pooled_output_markdown = self.activation(pooled_output_markdown)\n\n        outputs_code = self.roberta(\n            input_ids_code,\n            attention_mask=attention_mask_code,\n        )\n        sequence_output_code = outputs_code[0]\n\n        # pooled_output_code = sequence_output_code[:, 0, :]\n        pooled_output_code = (sequence_output_code * attention_mask_code.unsqueeze(-1)).sum(1) / attention_mask_code.unsqueeze(-1).sum(1)\n        pooled_output_code = self.dense(pooled_output_code)\n        pooled_output_code = self.activation(pooled_output_code)\n\n        logits = torch.mm(pooled_output_markdown, pooled_output_code.t())\n\n        loss = None\n        if labels is not None:\n            assert logits.shape[0] == logits.shape[1]  # square\n\n            loss_fct = CrossEntropyLoss()\n            loss = loss_fct(logits, labels)\n\n        if not return_dict:\n            output = (logits,) + outputs_markdown[2:]\n            return ((loss,) + output) if loss is not None else output\n\n        return SequenceClassifierOutput(\n            loss=loss,\n            logits=logits,\n            hidden_states=outputs_markdown.hidden_states,\n            attentions=outputs_markdown.attentions,\n        )\n\ndf_test[\"rank\"] = df_test.groupby([\"id\", \"cell_type\"]).cumcount()\ndf_test[\"pred\"] = df_test.groupby([\"id\", \"cell_type\"])[\"rank\"].rank(pct=False)\n\nall_preds = []\n\nfor model_path in model_paths:\n    tokenizer = AutoTokenizer.from_pretrained(model_path)   \n    model = MarkdownModelPairwise.from_pretrained(model_path)\n    \n    model.cuda()\n    model.eval()\n    \n    pred_ranks = []\n\n    for _, df_id in tqdm(df_test.groupby('id')):\n        df_id_markdown = df_id[df_id.cell_type == 'markdown']\n        df_id_code = df_id[df_id.cell_type != 'markdown']\n\n        df_id_code_rank = df_id_code['rank'].values\n        df_id_code_rank = df_id_code_rank * (df_id_code_rank[-1] + 1) / df_id_code_rank[-1]  # from 0 to n_code\n\n        n_markdown = len(df_id_markdown)\n        n_code = len(df_id_code_rank)\n\n        sources_markdown = df_id_markdown.source.astype(str).values\n        sources_code = df_id_code.source.astype(str).values\n\n        encodings_markdown = tokenizer(sources_markdown.tolist(), padding=True, truncation=True, max_length=max_length, return_tensors='pt')\n        encodings_code = tokenizer(sources_code.tolist(), padding=True, truncation=True, max_length=max_length, return_tensors='pt')\n\n        input_ids_markdown = encodings_markdown['input_ids'].cuda()\n        attention_mask_markdown = encodings_markdown['attention_mask'].cuda()\n        input_ids_code = encodings_code['input_ids'].cuda()\n        attention_mask_code = encodings_code['attention_mask'].cuda()\n\n        with torch.no_grad():\n            outputs = model(\n                input_ids=input_ids_markdown,\n                attention_mask=attention_mask_markdown,\n                input_ids_code=input_ids_code,\n                attention_mask_code=attention_mask_code,\n            )\n\n        logits = outputs.logits\n        assert logits.shape[0] == len(df_id_markdown)\n        assert logits.shape[1] == len(df_id_code)\n\n        preds = nn.functional.softmax(logits / temperature, dim=-1)\n\n        ranks = np.sum(preds.cpu().detach().numpy() * df_id_code_rank.reshape(1, -1), axis=-1)\n        pred_ranks = np.concatenate([pred_ranks, ranks/n_code])\n    \n    all_preds.append(pred_ranks)\n    \npairwise_1 = np.mean(all_preds, axis=0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"rank1","metadata":{}},{"cell_type":"code","source":"class MarkdownModelRank(nn.Module):\n    def __init__(self, model_path):\n        super(MarkdownModelRank, self).__init__()\n        self.model = AutoModel.from_pretrained(model_path)\n        self.top = nn.Linear(769, 1)\n        \n    def forward(self, ids, mask, fts):\n        x = self.model(ids, mask)[0]\n        x = self.top(torch.cat((x[:, 0, :], fts),1))\n        return x\n\nclass MarkdownDataset(Dataset):\n\n    def __init__(self, df, model_name_or_path, total_max_len, md_max_len, fts):\n        super().__init__()\n        self.df = df.reset_index(drop=True)\n        self.md_max_len = md_max_len\n        self.total_max_len = total_max_len  # maxlen allowed by model config\n        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n        self.fts = fts\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n\n        inputs = self.tokenizer.encode_plus(\n            row.source,\n            None,\n            add_special_tokens=True,\n            max_length=self.md_max_len,\n            padding=\"max_length\",\n            return_token_type_ids=True,\n            truncation=True\n        )\n        code_inputs = self.tokenizer.batch_encode_plus(\n            [str(x) for x in self.fts[row.id][\"codes\"]],\n            add_special_tokens=True,\n            max_length=15,\n            padding=\"max_length\",\n            truncation=True\n        )\n        n_md = self.fts[row.id][\"total_md\"]\n        n_code = self.fts[row.id][\"total_md\"]\n        if n_md + n_code == 0:\n            fts = torch.FloatTensor([0])\n        else:\n            fts = torch.FloatTensor([n_md / (n_md + n_code)])\n\n        ids = inputs['input_ids']\n        for x in code_inputs['input_ids']:\n            ids.extend(x[:-1])\n        ids = ids[:self.total_max_len]\n        if len(ids) != self.total_max_len:\n            ids = ids + [self.tokenizer.pad_token_id, ] * (self.total_max_len - len(ids))\n        ids = torch.LongTensor(ids)\n\n        mask = inputs['attention_mask']\n        for x in code_inputs['attention_mask']:\n            mask.extend(x[:-1])\n        mask = mask[:self.total_max_len]\n        if len(mask) != self.total_max_len:\n            mask = mask + [self.tokenizer.pad_token_id, ] * (self.total_max_len - len(mask))\n        mask = torch.LongTensor(mask)\n\n        assert len(ids) == self.total_max_len\n\n        return ids, mask, fts, torch.FloatTensor([row.pct_rank])\n\n    def __len__(self):\n        return self.df.shape[0]\n    \ndef read_data(data):\n    return tuple(d.cuda() for d in data[:-1]), data[-1].cuda()\n\n\ndef validate(model, val_loader):\n    model.eval()\n    \n    tbar = tqdm(val_loader, file=sys.stdout)\n    \n    preds = []\n    labels = []\n\n    with torch.no_grad():\n        for idx, data in enumerate(tbar):\n            inputs, target = read_data(data)\n\n            pred = model(*inputs)\n\n            preds.append(pred.detach().cpu().numpy().ravel())\n            labels.append(target.detach().cpu().numpy().ravel())\n    \n    return np.concatenate(labels), np.concatenate(preds)\n\n\ndef predict(model_path, ckpt_path, test_fts_):\n    model = MarkdownModelRank(model_path)\n    model = model.cuda()\n    model.eval()\n    model.load_state_dict(torch.load(ckpt_path))\n    BS = 32\n    NW = 8\n    MAX_LEN = 64\n    test_df[\"pct_rank\"] = 0\n    test_ds = MarkdownDataset(test_df[test_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True), md_max_len=64,total_max_len=512, model_name_or_path=model_path, fts=test_fts_)\n    test_loader = DataLoader(test_ds, batch_size=BS, shuffle=False, num_workers=NW,\n                              pin_memory=False, drop_last=False)\n    _, y_test = validate(model, test_loader)\n    return y_test\n\nmodel_path = \"../input/codebert-base/codebert-base/\"\nckpt_path = \"../input/multicodeinputs330codes/epoch_2_model.bin\"\nrank_1 = predict(model_path, ckpt_path, test_fts1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"rank2","metadata":{}},{"cell_type":"code","source":"rank_2 = predict(model_path, ckpt_path, test_fts2)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"ensemble","metadata":{}},{"cell_type":"code","source":"test_df.loc[test_df[\"cell_type\"] == \"markdown\", \"pred\"] = pairwise_1 * 0.3245044806 + rank_1 * 0.3377477597 + rank_2 * 0.3377477597","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 二段階目","metadata":{}},{"cell_type":"markdown","source":"pariwise1","metadata":{}},{"cell_type":"code","source":"import math\n\n\ndef calc_logits(tokenizer, model, markdown_source, code_source, max_length):\n    encodings_markdown = tokenizer(markdown_source.tolist(), padding=True, truncation=True, max_length=max_length, return_tensors='pt')\n    encodings_code = tokenizer(code_source.tolist(), padding=True, truncation=True, max_length=max_length, return_tensors='pt')\n\n    input_ids_markdown = encodings_markdown['input_ids'].cuda()\n    attention_mask_markdown = encodings_markdown['attention_mask'].cuda()\n    input_ids_code = encodings_code['input_ids'].cuda()\n    attention_mask_code = encodings_code['attention_mask'].cuda()\n\n    with torch.no_grad():\n        outputs = model(\n            input_ids=input_ids_markdown,\n            attention_mask=attention_mask_markdown,\n            input_ids_code=input_ids_code,\n            attention_mask_code=attention_mask_code,\n        )\n\n    logits = outputs.logits\n    return logits\n\n\nmodel_path = '../input/ai4code-pairwise-model/codebert-1epochKTKA-4epochKA'\n\ntokenizer = AutoTokenizer.from_pretrained(model_path)   \nmodel = MarkdownModelPairwise.from_pretrained(model_path)\nmodel.cuda()\nmodel.eval()\n\nmax_length = 128\nwindow_size = 45\ntemperature = 0.7\n\npairwise_1_df = None\n\nfor idx, sub_df in tqdm(test_df.groupby('id')):\n    code_sub_df = sub_df[sub_df.cell_type == 'code']\n    n_codes = len(code_sub_df)\n    if n_codes <= window_size:\n        if pairwise_1_df is None:\n            pairwise_1_df = sub_df.copy()\n        else:\n            pairwise_1_df = pd.concat([pairwise_1_df, sub_df])\n        continue\n\n    markdown_sub_df = sub_df[sub_df.cell_type == 'markdown']\n    n_markdowns = len(markdown_sub_df)\n\n    logits = calc_logits(tokenizer, model, markdown_sub_df.source, code_sub_df.source, max_length)\n    assert logits.shape[0] == n_markdowns\n    assert logits.shape[1] == n_codes\n\n    for m_idx, (_, row) in enumerate(markdown_sub_df.iterrows()):\n        m_rank = n_codes * row.pred\n        c_rank_left = min(\n            math.ceil(n_codes - window_size / 2),\n            max(math.ceil(m_rank - window_size / 2), 0),\n        )\n        c_rank_right = max(\n            math.ceil(window_size / 2),\n            min(math.ceil(m_rank + window_size / 2), n_codes),\n        )\n\n        m_logits = logits[m_idx, c_rank_left:c_rank_right]\n        pred = nn.functional.softmax(m_logits / temperature, dim=-1)\n\n        # c_ranks = np.arange(c_rank_left, c_rank_right)  # (c_rank_right - c_rank_left) elements and from c_rank_left to (c_rank_right - 1)\n        n_ranks = c_rank_right - c_rank_left\n        c_ranks = np.arange(n_ranks) * n_ranks / (n_ranks - 1) + (c_rank_left - 1)  # (c_rank_right - c_rank_left) elements but from (c_rank_left - 1) to (c_rank_right - 1)\n\n        m_pred_rank = np.sum(pred.cpu().detach().numpy() * c_ranks, axis=-1)\n        m_pred_pct_rank = m_pred_rank / n_codes\n\n        sub_df.loc[sub_df.cell_id == row.cell_id, \"pred\"] = m_pred_pct_rank\n\n    if pairwise_1_df is None:\n        pairwise_1_df = sub_df.copy()\n    else:\n        pairwise_1_df = pd.concat([pairwise_1_df, sub_df])\n\ndel tokenizer, model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"pairwise2","metadata":{}},{"cell_type":"code","source":"model_path = '../input/ai4code-pairwise-model/graphcodebert-1epochKTKA-4epochKA'\n\ntokenizer = AutoTokenizer.from_pretrained(model_path)   \nmodel = MarkdownModelPairwise.from_pretrained(model_path)\nmodel.cuda()\nmodel.eval()\n\nmax_length = 128\nwindow_size = 60\ntemperature = 0.7\n\npairwise_2_df = None\n\nfor idx, sub_df in tqdm(test_df.groupby('id')):\n    code_sub_df = sub_df[sub_df.cell_type == 'code']\n    n_codes = len(code_sub_df)\n    if n_codes <= window_size:\n        if pairwise_2_df is None:\n            pairwise_2_df = sub_df.copy()\n        else:\n            pairwise_2_df = pd.concat([pairwise_2_df, sub_df])\n        continue\n\n    markdown_sub_df = sub_df[sub_df.cell_type == 'markdown']\n    n_markdowns = len(markdown_sub_df)\n\n    logits = calc_logits(tokenizer, model, markdown_sub_df.source, code_sub_df.source, max_length)\n    assert logits.shape[0] == n_markdowns\n    assert logits.shape[1] == n_codes\n\n    for m_idx, (_, row) in enumerate(markdown_sub_df.iterrows()):\n        m_rank = n_codes * row.pred\n        c_rank_left = min(\n            math.ceil(n_codes - window_size / 2),\n            max(math.ceil(m_rank - window_size / 2), 0),\n        )\n        c_rank_right = max(\n            math.ceil(window_size / 2),\n            min(math.ceil(m_rank + window_size / 2), n_codes),\n        )\n\n        m_logits = logits[m_idx, c_rank_left:c_rank_right]\n        pred = nn.functional.softmax(m_logits / temperature, dim=-1)\n\n        # c_ranks = np.arange(c_rank_left, c_rank_right)  # (c_rank_right - c_rank_left) elements and from c_rank_left to (c_rank_right - 1)\n        n_ranks = c_rank_right - c_rank_left\n        c_ranks = np.arange(n_ranks) * n_ranks / (n_ranks - 1) + (c_rank_left - 1)  # (c_rank_right - c_rank_left) elements but from (c_rank_left - 1) to (c_rank_right - 1)\n\n        m_pred_rank = np.sum(pred.cpu().detach().numpy() * c_ranks, axis=-1)\n        m_pred_pct_rank = m_pred_rank / n_codes\n\n        sub_df.loc[sub_df.cell_id == row.cell_id, \"pred\"] = m_pred_pct_rank\n\n    if pairwise_2_df is None:\n        pairwise_2_df = sub_df.copy()\n    else:\n        pairwise_2_df = pd.concat([pairwise_2_df, sub_df])\n\ndel tokenizer, model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"rank1","metadata":{}},{"cell_type":"code","source":"class MarkdownDatasetFocus(Dataset):\n    def __init__(self, md, codes, model_name_or_path, total_max_len, md_max_len):\n        super().__init__()\n        self.md = md\n        self.codes = codes\n        self.md_max_len = md_max_len\n        self.total_max_len = total_max_len  # maxlen allowed by model config\n        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n\n    def __getitem__(self, index):\n        md = self.md[index]\n        codes = self.codes[index]\n\n        inputs = self.tokenizer.encode_plus(\n            md,\n            None,\n            add_special_tokens=True,\n            max_length=self.md_max_len,\n            padding=\"max_length\",\n            return_token_type_ids=True,\n            truncation=True,\n        )\n        code_inputs = self.tokenizer.batch_encode_plus(\n            [str(x) for x in codes],\n            add_special_tokens=True,\n            max_length=23,\n            padding=\"max_length\",\n            truncation=True,\n        )\n\n        fts = torch.FloatTensor([1 / (1 + len(codes))])\n\n        ids = inputs[\"input_ids\"]\n        for x in code_inputs[\"input_ids\"]:\n            ids.extend(x[:-1])\n        ids = ids[: self.total_max_len]\n        if len(ids) != self.total_max_len:\n            ids = ids + [\n                self.tokenizer.pad_token_id,\n            ] * (self.total_max_len - len(ids))\n        ids = torch.LongTensor(ids)\n\n        mask = inputs[\"attention_mask\"]\n        for x in code_inputs[\"attention_mask\"]:\n            mask.extend(x[:-1])\n        mask = mask[: self.total_max_len]\n        if len(mask) != self.total_max_len:\n            mask = mask + [\n                self.tokenizer.pad_token_id,\n            ] * (self.total_max_len - len(mask))\n        mask = torch.LongTensor(mask)\n\n        assert len(ids) == self.total_max_len\n\n        return ids, mask, fts, 0\n\n    def __len__(self):\n        return len(self.md)\n\n\nwindow_size = 20\n\nmodel_path = \"../input/codebert-base/codebert-base/\"\nckpt_path = \"../input/focus-ktka/20_codebert/20_codebert/model4.bin\"\n\nmodel = MarkdownModelRank(model_path)\nmodel = model.cuda()\nmodel.eval()\nmodel.load_state_dict(torch.load(ckpt_path))\n\nmd = []\ncode = []\n\nfor idx, sub_df in tqdm(test_df.groupby(\"id\")):\n    code_sub_df = sub_df[sub_df.cell_type == \"code\"]\n    n_codes = len(code_sub_df)\n    if n_codes <= window_size:\n        continue\n\n    markdown_sub_df = sub_df[sub_df.cell_type == \"markdown\"]\n    n_markdowns = len(markdown_sub_df)\n\n    for m_idx, (_, row) in enumerate(markdown_sub_df.iterrows()):\n        m_rank = n_codes * row.pred\n        c_rank_left = min(\n            math.ceil(n_codes - window_size / 2),\n            max(math.ceil(m_rank - window_size / 2), 0),\n        )\n        c_rank_right = max(\n            math.ceil(window_size / 2),\n            min(math.ceil(m_rank + window_size / 2), n_codes),\n        )\n\n        codes = code_sub_df[c_rank_left:c_rank_right][\"source\"].tolist()\n\n        md.append(row.source)\n        code.append(codes)\n\nds = MarkdownDatasetFocus(\n    md,\n    code,\n    md_max_len=64,\n    total_max_len=512,\n    model_name_or_path=model_path,\n)\nloader = DataLoader(\n    ds,\n    batch_size=32,\n    shuffle=False,\n    num_workers=8,\n    pin_memory=False,\n    drop_last=False,\n)\n\npreds = []\n\nwith torch.no_grad():\n    for data in tqdm(loader):\n        inputs, target = read_data(data)\n        pred = model(*inputs)\n        preds.append(pred.cpu().detach().numpy().ravel())\n\nif preds:\n    preds = np.concatenate(preds)\n\ncnt = 0\nrank_1_df = None\n\nfor idx, sub_df in tqdm(test_df.groupby(\"id\")):\n    \n    code_sub_df = sub_df[sub_df.cell_type == \"code\"]\n    n_codes = len(code_sub_df)\n    if n_codes <= window_size:\n        \n        if rank_1_df is None:\n            rank_1_df = sub_df.copy()\n        else:\n            rank_1_df = pd.concat([rank_1_df, sub_df])\n        \n        continue\n\n    markdown_sub_df = sub_df[sub_df.cell_type == \"markdown\"]\n\n    for m_idx, (_, row) in enumerate(markdown_sub_df.iterrows()):\n        m_rank = n_codes * row.pred\n        c_rank_left = min(\n            math.ceil(n_codes - window_size / 2),\n            max(math.ceil(m_rank - window_size / 2), 0),\n        )\n        c_rank_right = max(\n            math.ceil(window_size / 2),\n            min(math.ceil(m_rank + window_size / 2), n_codes),\n        )\n\n        pred = preds[cnt]\n        cnt += 1\n\n        c_pct_rank_left = code_sub_df.iloc[c_rank_left][\"pred\"]\n        c_pct_rank_right = code_sub_df.iloc[c_rank_right - 1][\"pred\"]\n\n        sub_df.loc[sub_df.cell_id == row.cell_id, \"pred\"] = (\n            c_pct_rank_left + (c_pct_rank_right - c_pct_rank_left) * pred\n        )\n\n    if rank_1_df is None:\n        rank_1_df = sub_df.copy()\n    else:\n        rank_1_df = pd.concat([rank_1_df, sub_df])\n\ndel md, code, model, ds, loader, preds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"rank2","metadata":{}},{"cell_type":"code","source":"# window_size = 30\n\n# model_path = \"../input/graphcodebert-base/graphcodebert-base/\"\n# ckpt_path = \"../input/focus-ktka/30_graphcodebert/30_graphcodebert/model4.bin\"\n\n# model = MarkdownModelRank(model_path)\n# model = model.cuda()\n# model.eval()\n# model.load_state_dict(torch.load(ckpt_path))\n\n# md = []\n# code = []\n\n# for idx, sub_df in tqdm(test_df.groupby(\"id\")):\n#     code_sub_df = sub_df[sub_df.cell_type == \"code\"]\n#     n_codes = len(code_sub_df)\n#     if n_codes <= window_size:\n#         continue\n\n#     markdown_sub_df = sub_df[sub_df.cell_type == \"markdown\"]\n#     n_markdowns = len(markdown_sub_df)\n\n#     for m_idx, (_, row) in enumerate(markdown_sub_df.iterrows()):\n#         m_rank = n_codes * row.pred\n#         c_rank_left = min(\n#             math.ceil(n_codes - window_size / 2),\n#             max(math.ceil(m_rank - window_size / 2), 0),\n#         )\n#         c_rank_right = max(\n#             math.ceil(window_size / 2),\n#             min(math.ceil(m_rank + window_size / 2), n_codes),\n#         )\n\n#         codes = code_sub_df[c_rank_left:c_rank_right][\"source\"].tolist()\n\n#         md.append(row.source)\n#         code.append(codes)\n\n# ds = MarkdownDatasetFocus(\n#     md,\n#     code,\n#     md_max_len=64,\n#     total_max_len=512,\n#     model_name_or_path=model_path,\n# )\n# loader = DataLoader(\n#     ds,\n#     batch_size=32,\n#     shuffle=False,\n#     num_workers=8,\n#     pin_memory=False,\n#     drop_last=False,\n# )\n\n# preds = []\n\n# with torch.no_grad():\n#     for data in tqdm(loader):\n#         inputs, target = read_data(data)\n#         pred = model(*inputs)\n#         preds.append(pred.cpu().detach().numpy().ravel())\n\n# if preds:\n#     preds = np.concatenate(preds)\n\n# cnt = 0\n# rank_2_df = None\n\n# for idx, sub_df in tqdm(test_df.groupby(\"id\")):\n    \n#     code_sub_df = sub_df[sub_df.cell_type == \"code\"]\n#     n_codes = len(code_sub_df)\n#     if n_codes <= window_size:\n        \n#         if rank_2_df is None:\n#             rank_2_df = sub_df.copy()\n#         else:\n#             rank_2_df = pd.concat([rank_2_df, sub_df])\n        \n#         continue\n\n#     markdown_sub_df = sub_df[sub_df.cell_type == \"markdown\"]\n\n#     for m_idx, (_, row) in enumerate(markdown_sub_df.iterrows()):\n#         m_rank = n_codes * row.pred\n#         c_rank_left = min(\n#             math.ceil(n_codes - window_size / 2),\n#             max(math.ceil(m_rank - window_size / 2), 0),\n#         )\n#         c_rank_right = max(\n#             math.ceil(window_size / 2),\n#             min(math.ceil(m_rank + window_size / 2), n_codes),\n#         )\n\n#         pred = preds[cnt]\n#         cnt += 1\n\n#         c_pct_rank_left = code_sub_df.iloc[c_rank_left][\"pred\"]\n#         c_pct_rank_right = code_sub_df.iloc[c_rank_right - 1][\"pred\"]\n\n#         sub_df.loc[sub_df.cell_id == row.cell_id, \"pred\"] = (\n#             c_pct_rank_left + (c_pct_rank_right - c_pct_rank_left) * pred\n#         )\n\n#     if rank_2_df is None:\n#         rank_2_df = sub_df.copy()\n#     else:\n#         rank_2_df = pd.concat([rank_2_df, sub_df])\n\n# del md, code, model, ds, loader, preds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"ensemble","metadata":{}},{"cell_type":"code","source":"new_df = pairwise_1_df.copy()\n\nnew_df.pred = pairwise_1_df.pred * 0.324555263 + pairwise_2_df.pred * 0.316037603 + rank_1_df.pred * 0.3594071339","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## submission","metadata":{}},{"cell_type":"code","source":"sub_df = new_df.sort_values(\"pred\").groupby(\"id\")[\"cell_id\"].apply(lambda x: \" \".join(x)).reset_index()\nsub_df.rename(columns={\"cell_id\": \"cell_order\"}, inplace=True)\nsub_df.to_csv(\"submission.csv\", index=False)\nsub_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}