Dopamin_evaluation.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GkMqGJ5SYISF",
        "outputId": "b9621c73-21f4-4c56-b8d1-d34999bee1a1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Cloning into 'code-comment-classification'...\n",
            "remote: Enumerating objects: 217, done.\u001b[K\n",
            "remote: Counting objects: 100% (56/56), done.\u001b[K\n",
            "remote: Compressing objects: 100% (37/37), done.\u001b[K\n",
            "remote: Total 217 (delta 28), reused 33 (delta 17), pack-reused 161\u001b[K\n",
            "Receiving objects: 100% (217/217), 20.29 MiB | 20.11 MiB/s, done.\n",
            "Resolving deltas: 100% (72/72), done.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/nlbse2024/code-comment-classification.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vw9QZiqQaFf7",
        "outputId": "a9b9fd56-ef14-4fb3-804e-f0a1a0cfb43b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting datasets\n",
            "  Downloading datasets-2.15.0-py3-none-any.whl (521 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m521.2/521.2 kB\u001b[0m \u001b[31m7.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.35.2)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.23.5)\n",
            "Requirement already satisfied: pyarrow>=8.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (9.0.0)\n",
            "Collecting pyarrow-hotfix (from datasets)\n",
            "  Downloading pyarrow_hotfix-0.6-py3-none-any.whl (7.9 kB)\n",
            "Collecting dill<0.3.8,>=0.3.0 (from datasets)\n",
            "  Downloading dill-0.3.7-py3-none-any.whl (115 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (1.5.3)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.31.0)\n",
            "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.1)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets) (3.4.1)\n",
            "Collecting multiprocess (from datasets)\n",
            "  Downloading multiprocess-0.70.15-py310-none-any.whl (134 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m10.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: fsspec[http]<=2023.10.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (2023.6.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.9.1)\n",
            "Requirement already satisfied: huggingface-hub>=0.18.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.19.4)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (23.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.13.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2023.6.3)\n",
            "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.15.0)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (23.1.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.9.3)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.4.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.18.0->datasets) (4.5.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->datasets) (2023.11.17)\n",
            "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2023.3.post1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.1->pandas->datasets) (1.16.0)\n",
            "Installing collected packages: pyarrow-hotfix, dill, multiprocess, datasets\n",
            "Successfully installed datasets-2.15.0 dill-0.3.7 multiprocess-0.70.15 pyarrow-hotfix-0.6\n"
          ]
        }
      ],
      "source": [
        "!pip install datasets transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "rK7qJNEsZvny"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "import time\n",
        "from datasets import Dataset\n",
        "from tqdm.auto import tqdm\n",
        "## Workaround for dashes in name\n",
        "from importlib import import_module\n",
        "nlbse_statistics = import_module('code-comment-classification.nlbse_statistics')\n",
        "from datasets import load_dataset, get_dataset_split_names\n",
        "from sklearn.metrics import f1_score, confusion_matrix\n",
        "import numpy as np\n",
        "tqdm.pandas()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XmSQgUdwba3I"
      },
      "source": [
        "# Model declaration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "gV1r0tF4bdkF"
      },
      "outputs": [],
      "source": [
        "from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaModel, RobertaClassificationHead, RobertaLayer\n",
        "from transformers.modeling_outputs import SequenceClassifierOutput\n",
        "import math\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss\n",
        "import torch.nn.init as init\n",
        "from typing import *\n",
        "\n",
        "class CustomRobertaClassificationHead(nn.Module):\n",
        "    \"\"\"Head for sentence-level classification tasks.\"\"\"\n",
        "\n",
        "    def __init__(self, config, infeat):\n",
        "        super().__init__()\n",
        "        self.linear = nn.Linear(infeat, config.hidden_size)\n",
        "        classifier_dropout = (\n",
        "            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob\n",
        "        )\n",
        "        self.dropout = nn.Dropout(classifier_dropout)\n",
        "        self.out = nn.Linear(config.hidden_size, config.num_labels)\n",
        "\n",
        "    def forward(self, features, **kwargs):\n",
        "        x = self.dropout(features)\n",
        "        x = self.linear(x)\n",
        "        x = torch.tanh(x)\n",
        "        x = self.dropout(x)\n",
        "        x = self.out(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class CodeBERTForSequenceClassification(RobertaPreTrainedModel):\n",
        "    def __init__(self, config):\n",
        "        super().__init__(config)\n",
        "        self.num_labels = config.num_labels\n",
        "        self.config = config\n",
        "\n",
        "        self.roberta = RobertaModel(config, add_pooling_layer=False)\n",
        "        self.classifier = CustomRobertaClassificationHead(config, 768)\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",
        "        weighted_loss = None\n",
        "    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:\n",
        "        r\"\"\"\n",
        "        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):\n",
        "            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,\n",
        "            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If\n",
        "            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).\n",
        "        \"\"\"\n",
        "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
        "\n",
        "        outputs = 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 = outputs[0][:,0,:]\n",
        "\n",
        "        logits = self.classifier(sequence_output)\n",
        "\n",
        "        loss = None\n",
        "        if labels is not None:\n",
        "            # move labels to correct device to enable model parallelism\n",
        "            labels = labels.to(logits.device)\n",
        "            if self.config.problem_type is None:\n",
        "                if self.num_labels == 1:\n",
        "                    self.config.problem_type = \"regression\"\n",
        "                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):\n",
        "                    self.config.problem_type = \"single_label_classification\"\n",
        "                else:\n",
        "                    self.config.problem_type = \"multi_label_classification\"\n",
        "\n",
        "            if self.config.problem_type == \"regression\":\n",
        "                loss_fct = MSELoss()\n",
        "                if self.num_labels == 1:\n",
        "                    loss = loss_fct(logits.squeeze(), labels.squeeze())\n",
        "                else:\n",
        "                    loss = loss_fct(logits, labels)\n",
        "            elif self.config.problem_type == \"single_label_classification\":\n",
        "                if weighted_loss is not None:\n",
        "                    loss_fct = CrossEntropyLoss(weight=weighted_loss[0][:self.num_labels])\n",
        "                else:\n",
        "                    loss_fct = CrossEntropyLoss()\n",
        "\n",
        "                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))\n",
        "            elif self.config.problem_type == \"multi_label_classification\":\n",
        "                loss_fct = BCEWithLogitsLoss()\n",
        "                loss = loss_fct(logits, labels)\n",
        "\n",
        "        if not return_dict:\n",
        "            output = (logits,) + outputs[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.hidden_states,\n",
        "            attentions=outputs.attentions,\n",
        "        )\n",
        "\n",
        "\n",
        "class HSUM(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super(HSUM, self).__init__()\n",
        "        self.config = config\n",
        "\n",
        "        self.pre_layers = torch.nn.ModuleList()\n",
        "        for i in range(config.count):\n",
        "            self.pre_layers.append(RobertaLayer(config))\n",
        "\n",
        "        self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        \"\"\"Initialize the weights\"\"\"\n",
        "        if isinstance(module, nn.Linear):\n",
        "            # Slightly different from the TF version which uses truncated_normal for initialization\n",
        "            # cf https://github.com/pytorch/pytorch/pull/5617\n",
        "            module.weight.data.normal_(mean=0.0, std=self.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.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 forward(self, layers, attention_mask, return_list= False):\n",
        "        logitses = []\n",
        "        output = torch.zeros_like(layers[0])\n",
        "\n",
        "        for i in range(self.config.count):\n",
        "            output = output + layers[-i-1]\n",
        "            logits = self.pre_layers[i](output, attention_mask)[0]\n",
        "            logitses.append(logits)\n",
        "\n",
        "        if return_list:\n",
        "            return logitses\n",
        "\n",
        "        avg_logits = torch.sum(torch.stack(logitses), dim=0) / self.config.count\n",
        "        return avg_logits\n",
        "\n",
        "\n",
        "class CodeBERTHSUMForSequenceClassification(RobertaPreTrainedModel):\n",
        "    def __init__(self, config):\n",
        "        super().__init__(config)\n",
        "        self.num_labels = config.num_labels\n",
        "        self.config = config\n",
        "\n",
        "        self.roberta = RobertaModel(config, add_pooling_layer=False)\n",
        "        self.classifier = CustomRobertaClassificationHead(config, 768)\n",
        "\n",
        "        self.mixlayer = HSUM(config)\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",
        "        weighted_loss = None\n",
        "    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:\n",
        "        r\"\"\"\n",
        "        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):\n",
        "            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,\n",
        "            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If\n",
        "            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).\n",
        "        \"\"\"\n",
        "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
        "        # if labels is not None:\n",
        "        output_hidden_states = True\n",
        "\n",
        "        outputs = 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",
        "\n",
        "        if return_dict:\n",
        "            layers = outputs.hidden_states\n",
        "        else:\n",
        "            layers = outputs[2]\n",
        "        extend_attention_mask = (1.0 - attention_mask[:,None, None, :]) * -10000.0\n",
        "\n",
        "        sequence_output = self.mixlayer(layers, extend_attention_mask)[:,0,:]\n",
        "\n",
        "        logits = self.classifier(sequence_output)\n",
        "\n",
        "        loss = None\n",
        "        if labels is not None:\n",
        "            # move labels to correct device to enable model parallelism\n",
        "            labels = labels.to(logits.device)\n",
        "            if self.config.problem_type is None:\n",
        "                if self.num_labels == 1:\n",
        "                    self.config.problem_type = \"regression\"\n",
        "                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):\n",
        "                    self.config.problem_type = \"single_label_classification\"\n",
        "                else:\n",
        "                    self.config.problem_type = \"multi_label_classification\"\n",
        "\n",
        "            if self.config.problem_type == \"regression\":\n",
        "                loss_fct = MSELoss()\n",
        "                if self.num_labels == 1:\n",
        "                    loss = loss_fct(logits.squeeze(), labels.squeeze())\n",
        "                else:\n",
        "                    loss = loss_fct(logits, labels)\n",
        "            elif self.config.problem_type == \"single_label_classification\":\n",
        "                if weighted_loss is not None:\n",
        "                    loss_fct = CrossEntropyLoss(weight=weighted_loss[0][:self.num_labels])\n",
        "                else:\n",
        "                    loss_fct = CrossEntropyLoss()\n",
        "\n",
        "                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))\n",
        "            elif self.config.problem_type == \"multi_label_classification\":\n",
        "                loss_fct = BCEWithLogitsLoss()\n",
        "                loss = loss_fct(logits, labels)\n",
        "\n",
        "        if not return_dict:\n",
        "            output = (logits,) + outputs[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.hidden_states,\n",
        "            attentions=outputs.attentions,\n",
        "        )\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AzMAyf1FhXl-"
      },
      "source": [
        "# Evaluation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "re8X3P3dmmfa"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "24iHC9xZaWz8"
      },
      "outputs": [],
      "source": [
        "langs = ['java', 'python', 'pharo']\n",
        "lan_cats = []\n",
        "datasets = {}\n",
        "for lan in langs: # for each language\n",
        "    df = pd.read_csv(f'./code-comment-classification/{lan}/input/{lan}.csv')\n",
        "    df['combo'] = df[['class', 'comment_sentence']].agg('</s>'.join, axis=1)\n",
        "    df['label'] = df.instance_type\n",
        "    df[\"combo_len\"] = [len(x.split()) for x in df[\"combo\"]]\n",
        "    cats = list(map(lambda x: lan + '_' + x, list(set(df.category))))\n",
        "\n",
        "    for cat in list(set(df.category)): # for each category\n",
        "        filtered =  df[df.category == cat].sort_values(by='combo_len')\n",
        "        train_data = Dataset.from_pandas(filtered[filtered.partition == 0])\n",
        "        test_data = Dataset.from_pandas(filtered[filtered.partition == 1])\n",
        "        datasets[f'{lan}_{cat}'.lower()] = {'train_data': train_data, 'test_data' : test_data}\n",
        "        lan_cats.append(f'{lan}_{cat}'.lower())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "x9BDyAz_Iol6"
      },
      "outputs": [],
      "source": [
        "def get_prediction(x, model, tokenizer, num_iter):\n",
        "  y_hat = []\n",
        "  for i in range(int(num_iter)):\n",
        "      inputs = tokenizer(x[i*batch_size:(i+1)*batch_size], max_length=32, padding=True, truncation=True, return_tensors=\"pt\")\n",
        "      logits = model(**inputs).logits\n",
        "      y_hat.extend(np.argmax(logits.cpu().numpy(), axis=1).tolist())\n",
        "  return y_hat"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "5QdGX1FKhnNI"
      },
      "outputs": [],
      "source": [
        "scores = []\n",
        "batch_size = 64\n",
        "torch.set_default_device('cuda')\n",
        "for lan_cat in lan_cats:\n",
        "    # load models and data\n",
        "    tokenizer = AutoTokenizer.from_pretrained(\"Fsoft-AIC/dopamin-{}\".format(lan_cat.replace(\"_\",\"-\")))\n",
        "    model = CodeBERTHSUMForSequenceClassification.from_pretrained(\"Fsoft-AIC/dopamin-{}\".format(lan_cat.replace(\"_\",\"-\")))\n",
        "    model.eval()\n",
        "    test_data = datasets[lan_cat]['test_data']\n",
        "    x = test_data[\"combo\"]\n",
        "    y = test_data['label']\n",
        "    num_iter = np.ceil(len(x) / batch_size)\n",
        "\n",
        "    # # run and time 10 times for each cat\n",
        "    with torch.no_grad():\n",
        "      for it in range(10):\n",
        "        ############# TIME BLOCK #####################\n",
        "        num_iter = np.ceil(len(x) / batch_size)\n",
        "        start = time.time()\n",
        "        y_hat = get_prediction(x, model, tokenizer, num_iter)\n",
        "        elapsed_time = time.time() - start\n",
        "        time_per_sample = elapsed_time / len(y)\n",
        "\n",
        "        # # ############# TIME BLOCK #####################\n",
        "        _, fp, fn, tp = confusion_matrix(y_hat, y).ravel()\n",
        "        wf1 = f1_score(y, y_hat, average='weighted')\n",
        "        precision, recall, f1 = nlbse_statistics.get_precision_recall_f1(tp, fp, fn)\n",
        "        scores.append({'lan_cat': lan_cat.lower(),'precision': precision,'recall': recall,'f1': f1,'wf1': wf1, 'avg_runtime': time_per_sample, 'iteration': it, 'len': len(test_data)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 711
        },
        "id": "mn5cgHmvnptX",
        "outputId": "1e601a08-f97c-482a-8693-786940481f32"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Average runtime: 0.00203\n",
            "Average f1: 0.74\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <div id=\"df-e0777f56-af1b-4b59-82de-34d4e673896a\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>precision</th>\n",
              "      <th>recall</th>\n",
              "      <th>f1</th>\n",
              "      <th>avg_runtime</th>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>lan_cat</th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>java_deprecation</th>\n",
              "      <td>0.81</td>\n",
              "      <td>0.88</td>\n",
              "      <td>0.8451</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>java_expand</th>\n",
              "      <td>0.52</td>\n",
              "      <td>0.63</td>\n",
              "      <td>0.5677</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>java_ownership</th>\n",
              "      <td>1.00</td>\n",
              "      <td>0.99</td>\n",
              "      <td>0.9956</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>java_pointer</th>\n",
              "      <td>0.89</td>\n",
              "      <td>0.81</td>\n",
              "      <td>0.8488</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>java_rational</th>\n",
              "      <td>0.49</td>\n",
              "      <td>0.59</td>\n",
              "      <td>0.5361</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>java_summary</th>\n",
              "      <td>0.94</td>\n",
              "      <td>0.87</td>\n",
              "      <td>0.9031</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>java_usage</th>\n",
              "      <td>0.87</td>\n",
              "      <td>0.95</td>\n",
              "      <td>0.9073</td>\n",
              "      <td>0.0023</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_classreferences</th>\n",
              "      <td>0.76</td>\n",
              "      <td>0.62</td>\n",
              "      <td>0.6842</td>\n",
              "      <td>0.0021</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_collaborators</th>\n",
              "      <td>0.57</td>\n",
              "      <td>0.64</td>\n",
              "      <td>0.6038</td>\n",
              "      <td>0.0021</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_example</th>\n",
              "      <td>0.93</td>\n",
              "      <td>0.93</td>\n",
              "      <td>0.9342</td>\n",
              "      <td>0.0020</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_intent</th>\n",
              "      <td>0.87</td>\n",
              "      <td>0.85</td>\n",
              "      <td>0.8571</td>\n",
              "      <td>0.0021</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_keyimplementationpoints</th>\n",
              "      <td>0.50</td>\n",
              "      <td>0.69</td>\n",
              "      <td>0.5783</td>\n",
              "      <td>0.0021</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_keymessages</th>\n",
              "      <td>0.86</td>\n",
              "      <td>0.81</td>\n",
              "      <td>0.8308</td>\n",
              "      <td>0.0019</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>pharo_responsibilities</th>\n",
              "      <td>0.59</td>\n",
              "      <td>0.62</td>\n",
              "      <td>0.6074</td>\n",
              "      <td>0.0022</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>python_developmentnotes</th>\n",
              "      <td>0.45</td>\n",
              "      <td>0.44</td>\n",
              "      <td>0.4427</td>\n",
              "      <td>0.0018</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>python_expand</th>\n",
              "      <td>0.51</td>\n",
              "      <td>0.53</td>\n",
              "      <td>0.5174</td>\n",
              "      <td>0.0018</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>python_parameters</th>\n",
              "      <td>0.88</td>\n",
              "      <td>0.85</td>\n",
              "      <td>0.8632</td>\n",
              "      <td>0.0017</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>python_summary</th>\n",
              "      <td>0.80</td>\n",
              "      <td>0.69</td>\n",
              "      <td>0.7400</td>\n",
              "      <td>0.0017</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>python_usage</th>\n",
              "      <td>0.72</td>\n",
              "      <td>0.79</td>\n",
              "      <td>0.7564</td>\n",
              "      <td>0.0018</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-e0777f56-af1b-4b59-82de-34d4e673896a')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-e0777f56-af1b-4b59-82de-34d4e673896a button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-e0777f56-af1b-4b59-82de-34d4e673896a');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-e2aeffc7-f854-44b8-b73f-be4f41bfcd65\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-e2aeffc7-f854-44b8-b73f-be4f41bfcd65')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-e2aeffc7-f854-44b8-b73f-be4f41bfcd65 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "text/plain": [
              "                               precision  recall      f1  avg_runtime\n",
              "lan_cat                                                              \n",
              "java_deprecation                    0.81    0.88  0.8451       0.0022\n",
              "java_expand                         0.52    0.63  0.5677       0.0022\n",
              "java_ownership                      1.00    0.99  0.9956       0.0022\n",
              "java_pointer                        0.89    0.81  0.8488       0.0022\n",
              "java_rational                       0.49    0.59  0.5361       0.0022\n",
              "java_summary                        0.94    0.87  0.9031       0.0022\n",
              "java_usage                          0.87    0.95  0.9073       0.0023\n",
              "pharo_classreferences               0.76    0.62  0.6842       0.0021\n",
              "pharo_collaborators                 0.57    0.64  0.6038       0.0021\n",
              "pharo_example                       0.93    0.93  0.9342       0.0020\n",
              "pharo_intent                        0.87    0.85  0.8571       0.0021\n",
              "pharo_keyimplementationpoints       0.50    0.69  0.5783       0.0021\n",
              "pharo_keymessages                   0.86    0.81  0.8308       0.0019\n",
              "pharo_responsibilities              0.59    0.62  0.6074       0.0022\n",
              "python_developmentnotes             0.45    0.44  0.4427       0.0018\n",
              "python_expand                       0.51    0.53  0.5174       0.0018\n",
              "python_parameters                   0.88    0.85  0.8632       0.0017\n",
              "python_summary                      0.80    0.69  0.7400       0.0017\n",
              "python_usage                        0.72    0.79  0.7564       0.0018"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df = pd.DataFrame(scores).groupby('lan_cat').mean()\n",
        "df['time_std'] = pd.DataFrame(scores).groupby('lan_cat').std()['avg_runtime']\n",
        "avg_runtime = round(df['avg_runtime'].mean(), 5)\n",
        "avg_score = round(df['f1'].mean(), 2)\n",
        "print(f\"Average runtime: {round(df['avg_runtime'].mean(), 5)}\")\n",
        "print(f\"Average f1: {round(df['f1'].mean(), 2)}\")\n",
        "df.precision = df.precision.round(2)\n",
        "df.recall = df.recall.round(2)\n",
        "df.f1 = df.f1.round(4)\n",
        "df.avg_runtime = df.avg_runtime.round(4)\n",
        "\n",
        "df[['precision','recall', 'f1', 'avg_runtime']]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rXlyFiEmyY5U",
        "outputId": "819afdc1-1829-4008-edfc-0494317139fe"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0.7034999999999999"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "def score(f1, avg_time_per_sample):\n",
        "  return 0.75 * f1 + 0.25 * ((max((0.005 - avg_time_per_sample), 0)/0.005))\n",
        "score(avg_score,avg_runtime)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}