{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5da2447e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# !pip install gdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3ec45c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import SATD_tool\n",
    "import importlib\n",
    "import gzip\n",
    "importlib.reload(SATD_tool)\n",
    "from SATD_tool import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fa0e0c5",
   "metadata": {},
   "source": [
    "# Download Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c99a2c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gdown\n",
    "\n",
    "def download_from_google_drive(fileId, filePath):\n",
    "    download_url = f\"https://drive.google.com/uc?id={fileId}\"\n",
    "    gdown.download(download_url, filePath, quiet=False)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "42a66f26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datasets/Python/df3.pkl.gz already exists.\n",
      "Datasets/Python/df3_train.pkl.gz already exists.\n",
      "Datasets/Python/df3_test.pkl.gz already exists.\n",
      "Datasets/Python/df3_dev.pkl.gz already exists.\n"
     ]
    }
   ],
   "source": [
    "# Download Python dataset\n",
    "\n",
    "os.makedirs(\"Datasets/Python\", exist_ok=True)\n",
    "\n",
    "fileId_filePath = {\n",
    "    '1aQ5Rpe-GI-Vqb5BCg8cz8bXXEKzKmguX': 'Datasets/Python/df3.pkl.gz',\n",
    "    '1obp8vk7wBSguKy72UPPCLHOcVRTGYGuP': 'Datasets/Python/df3_train.pkl.gz',\n",
    "    '1ItMGQfbZHMbnvhljbJqr2dP4WFLbjQky': 'Datasets/Python/df3_test.pkl.gz',\n",
    "    '1ePmAuajVjYiKfNKLPhNsCIp9XK0oTx-Q': 'Datasets/Python/df3_dev.pkl.gz',\n",
    "}\n",
    "\n",
    "for fileId, filePath in fileId_filePath.items():\n",
    "    if not os.path.exists(filePath):\n",
    "        download_from_google_drive(fileId, filePath)\n",
    "    else:\n",
    "        print(filePath, 'already exists.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1408538e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datasets/Java/df3.pkl.gz already exists.\n",
      "Datasets/Java/df3_train.pkl.gz already exists.\n",
      "Datasets/Java/df3_test.pkl.gz already exists.\n",
      "Datasets/Java/df3_dev.pkl.gz already exists.\n"
     ]
    }
   ],
   "source": [
    "# Download Java dataset\n",
    "\n",
    "os.makedirs(\"Datasets/Java\", exist_ok=True)\n",
    "\n",
    "fileId_filePath = {\n",
    "    '1t5Pf0f8NSygdNBgTxtdsPPmvGbmFlbFe': 'Datasets/Java/df3.pkl.gz',\n",
    "    '16K6kc_9_sOsFDw9z2sOp-33AE1BREe2P': 'Datasets/Java/df3_train.pkl.gz',\n",
    "    '1ecb3uLQ-DeT1FgI0p8Yxk6SCfqluTkLI': 'Datasets/Java/df3_test.pkl.gz',\n",
    "    '1KJq4RNupQW8GH9wuE48tbZtHOCCbFGrv': 'Datasets/Java/df3_dev.pkl.gz',\n",
    "}\n",
    "\n",
    "for fileId, filePath in fileId_filePath.items():\n",
    "    if not os.path.exists(filePath):\n",
    "        download_from_google_drive(fileId, filePath)\n",
    "    else:\n",
    "        print(filePath, 'already exists.')        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2c56bb39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datasets/Mastropaolo/mastropaolo_with_filter_columns.pkl.gz already exists.\n"
     ]
    }
   ],
   "source": [
    "# Download Mastropaolo dataset\n",
    "\n",
    "os.makedirs(\"Datasets/Mastropaolo\", exist_ok=True)\n",
    "\n",
    "fileId_filePath = {\n",
    "    '1PgeYVvyk1yu0s8AxEcFrK5CjBlUPvKAq': 'Datasets/Mastropaolo/mastropaolo_with_filter_columns.pkl.gz',\n",
    "}\n",
    "\n",
    "for fileId, filePath in fileId_filePath.items():\n",
    "    if not os.path.exists(filePath):\n",
    "        download_from_google_drive(fileId, filePath)\n",
    "    else:\n",
    "        print(filePath, 'already exists.')    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4963b0d",
   "metadata": {},
   "source": [
    "# Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "25b1e389",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DATASET: Java, LANGUAGE:Java\n"
     ]
    }
   ],
   "source": [
    "# RQ2: uses the first 1000 items in the dataset for evaluating models by prompt engineering.\n",
    "# Note that each item in the dataset has a randomly assinged index (rand_index). As the dataset is now sorted by the rand_index,\n",
    "# taking the first 1000 rows means that we randomly select 1000 items from the dataset.\n",
    "\n",
    "# RQ3: uses the Mastropaolo dataset\n",
    "\n",
    "# RQ4: uses the whole dataset which is repository-based splitted to train, dev, and test.\n",
    "\n",
    "RQ = 'RQ2' # RQ2 or RQ3 or RQ4 \n",
    "LANGUAGE = 'Java' # Python or Java (no need to set for RQ3)\n",
    "\n",
    "if RQ=='RQ2':\n",
    "    DATASET = LANGUAGE\n",
    "elif RQ=='RQ3':\n",
    "    LANGUAGE = 'Java'\n",
    "    DATASET = 'Mastropaolo'\n",
    "elif RQ=='RQ4':\n",
    "    DATASET = LANGUAGE + '_test'\n",
    "else:\n",
    "    raise ValueError(\"Invalid value for RQ:\", RQ)\n",
    "\n",
    "print(f\"DATASET: {DATASET}, LANGUAGE:{LANGUAGE}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e47d994d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1000, 26)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Index(['rand_index', 'user', 'project', 'created_in_file',\n",
       "       'last_appeared_in_file', 'created_in_line', 'last_appeared_in_line',\n",
       "       'created_in_commit', 'deleted_in_commit', 'created_at_date',\n",
       "       'deleted_at_date', 'content', 'deleted_in_lines', 'created_in_lines',\n",
       "       'updated_in_commits', 'last_content', 'SATD_comment',\n",
       "       'containing_method_applied_approach',\n",
       "       'containing_method_before_repayment',\n",
       "       'containing_method_after_repayment', 'method_is_updated',\n",
       "       'SATD_count_before_repayment', 'SATD_count_after_repayment',\n",
       "       'method_tokens_before_repayment', 'method_tokens_after_repayment',\n",
       "       'is_repayment_llama3'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "DATASETS_DIR = 'Datasets'\n",
    "\n",
    "if RQ=='RQ2':\n",
    "    df = pd.read_pickle(gzip.open(f'{DATASETS_DIR}/{LANGUAGE}/df3.pkl.gz', 'rb'))\n",
    "    df = df[df['is_repayment_llama3'] == 'yes'] # we need to apply the last filter\n",
    "    df = df.head(1000)\n",
    "elif RQ=='RQ3':\n",
    "    df = pd.read_pickle(gzip.open(f'{DATASETS_DIR}/Mastropaolo/mastropaolo_with_filter_columns.pkl.gz', 'rb'))\n",
    "    df = df[df['data_split']=='test']\n",
    "    # apply our two new filters\n",
    "    df = df[(df[\"SATD_count_before_repayment\"] == 1) & (df[\"SATD_count_after_repayment\"] == 0)]\n",
    "    df = df[df['is_repayment_llama3']=='yes']   \n",
    "elif RQ=='RQ4':\n",
    "    df = pd.read_pickle(gzip.open(f'{DATASETS_DIR}/{LANGUAGE}/df3_test.pkl.gz', 'rb'))\n",
    "    df_train = pd.read_pickle(f'/home/jovyan/SATD-Repayment/Create Dataset/Datasets/{LANGUAGE}/df3_train.pkl')\n",
    "print(df.shape)\n",
    "df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "699a0935",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add deletion_commit_url\n",
    "def get_deletion_commit_url(row):\n",
    "    return f\"\"\"www.github.com/{row['user']}/{row['project']}/commit/{row['deleted_in_commit']}\"\"\"\n",
    "\n",
    "if RQ=='RQ2':\n",
    "    df['deletion_commit_url'] = df.apply(get_deletion_commit_url, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7d2d5efa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of items that 2 or less new lines are added or updated (Easy group): 492 (%49.2)\n",
      "Number of items that 3 or more new lines are added or updated (Hard group): 508 (%50.8)\n"
     ]
    }
   ],
   "source": [
    "def get_number_of_deleted_and_inserted_lines(row):\n",
    "    deleted_lines, inserted_lines = get_deleted_and_inserted_lines(row['containing_method_before_repayment'], row['containing_method_after_repayment'])\n",
    "    return len(deleted_lines), len(inserted_lines)\n",
    "\n",
    "df[['number_of_deleted_lines', 'number_of_inserted_lines']] = df.apply(\n",
    "    lambda row: pd.Series(get_number_of_deleted_and_inserted_lines(row)), axis=1\n",
    ")\n",
    "\n",
    "i=2\n",
    "count = len(df[df['number_of_inserted_lines']<=i])\n",
    "print(f'Number of items that {i} or less new lines are added or updated (Easy group): {count} (%{100*count/len(df):.1f})')\n",
    "count = len(df[df['number_of_inserted_lines']>i])\n",
    "print(f'Number of items that {i+1} or more new lines are added or updated (Hard group): {count} (%{100*count/len(df):.1f})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9cab0ac",
   "metadata": {},
   "source": [
    "**Add the prompt column to the dataframe**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4cecbc9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROMPT_TEMPLATE: NoExplain\n",
      "\n",
      "How to update the following code to resolve the SATD? No need to explain. Just provide the updated code.\n",
      "\n",
      "### Code:\n",
      "public Connection readConnection(URI connectionURI) throws NoSuchConnectionException {\n",
      "    logger.debug(MessageFormat.format(\"need-facing: READ_CONNECTION called for connection {0}\", connectionURI));\n",
      "\n",
      "    //TODO: Convert readConnectionContent(connectionURI)\n",
      "    return null;\n",
      "}\n",
      "\n",
      "### SATD comment:\n",
      "TODO: Convert readConnectionContent(connectionURI)\n",
      "\n",
      "### Updated code after SATD repayment:\n"
     ]
    }
   ],
   "source": [
    "PROMPT_TEMPLATE = 'NoExplain' # 'Mastropaolo-T2' or 'NoExplain' or 'CoT1' or 'CoT2'\n",
    "\n",
    "def get_prompt(row):\n",
    "    \n",
    "    if PROMPT_TEMPLATE == 'Mastropaolo-T2':\n",
    "        prompt = f\"\"\"Perform removal of this SATD: \"{row['SATD_comment']}\" from this code:\n",
    "\n",
    "{row['containing_method_before_repayment']}\n",
    "\"\"\"\n",
    "    \n",
    "    if PROMPT_TEMPLATE == 'NoExplain':\n",
    "        prompt = f\"\"\"How to update the following code to resolve the SATD? No need to explain. Just provide the updated code.\n",
    "\n",
    "### Code:\n",
    "{row['containing_method_before_repayment']}\n",
    "\n",
    "### SATD comment:\n",
    "{row['SATD_comment']}\n",
    "\n",
    "### Updated code after SATD repayment:\"\"\"\n",
    "\n",
    "    if PROMPT_TEMPLATE == 'CoT1':\n",
    "        prompt = f\"\"\"How to update the following code to resolve the SATD?\n",
    "\n",
    "### Code:\n",
    "{row['containing_method_before_repayment']}\n",
    "\n",
    "### SATD comment:\n",
    "{row['SATD_comment']}\n",
    "\n",
    "### Consider the following questions in your answer:\n",
    "Shortly explain how to resolve the SATD.\n",
    "Provide the updated code.\"\"\"\n",
    "    \n",
    "    if PROMPT_TEMPLATE == 'CoT2':\n",
    "        prompt = f\"\"\"How can the following code be updated to resolve the Self-Admitted Technical Debt (SATD)?\n",
    "\n",
    "### Code:\n",
    "{row['containing_method_before_repayment']}\n",
    "\n",
    "### SATD comment:\n",
    "{row['SATD_comment']}\n",
    "\n",
    "### Consider the following questions in your answer:\n",
    "1. Briefly explain how to resolve the SATD.\n",
    "2. Provide the updated code.\"\"\"\n",
    "    \n",
    "    return prompt\n",
    "\n",
    "\n",
    "df['prompt2'] = df.apply(get_prompt, axis=1)\n",
    "print(f'PROMPT_TEMPLATE: {PROMPT_TEMPLATE}\\n')\n",
    "print(df.iloc[0]['prompt2'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a90f5c93",
   "metadata": {},
   "source": [
    "# Use OpenAI - GPT4o-mini\n",
    "\n",
    "This part is to reproduce the generated outputs by GPT4o-mini. You can pass to the next section if you don't want to reproduce them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f120e236",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "client = OpenAI(api_key=\"xxx\") # provide your OpenAI API key\n",
    "\n",
    "# test it\n",
    "response = client.chat.completions.create(\n",
    "  model=\"gpt-4o-mini\",\n",
    "  temperature=0.0,\n",
    "  messages=[\n",
    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Who won the world series in 2020?\"}\n",
    "  ]\n",
    ")\n",
    "response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "118b3825",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "def create_batch_file(df, start_row, end_row, prompt_template, openai_model, dataset):\n",
    "    filename = dataset + '_' + prompt_template + '_' + openai_model + '_batch' + str(start_row) + 'to' + str(end_row) + '.jsonl'\n",
    "    with open(filename, 'w') as f:\n",
    "        for i in range(start_row, end_row):\n",
    "            batch = {}\n",
    "            batch['custom_id'] = str(df.iloc[i]['rand_index'])\n",
    "            batch['method'] = 'POST'\n",
    "            batch['url'] = '/v1/chat/completions'\n",
    "            body = {\"model\": openai_model, 'temperature':0.0, \"messages\": [{\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},{\"role\": \"user\", \"content\": df.iloc[i]['prompt2']}],\"max_tokens\": 2048}\n",
    "            batch['body'] = body\n",
    "            f.write(json.dumps(batch) + '\\n')\n",
    "    return filename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4467aa6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch file saved in Java_test_Mastropaolo-T2_gpt-4o-mini_batch0to9013.jsonl\n"
     ]
    }
   ],
   "source": [
    "# create batch file\n",
    "openai_model = \"gpt-4o-mini\" # gpt-4o-mini (current: gpt-4o-mini-2024-07-18)\n",
    "filename = create_batch_file(df, 0, len(df), PROMPT_TEMPLATE, openai_model, DATASET)\n",
    "print('Batch file saved in', filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1ea3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload the file to OpenAI\n",
    "batch_input_file = client.files.create(\n",
    "  file=open(filename, \"rb\"),\n",
    "  purpose=\"batch\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae0c280",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the batch (it starts the process)\n",
    "# NOTE: make sure you don't run it twice, or you will be charged twice\n",
    "\n",
    "batch_input_file_id = batch_input_file.id\n",
    "client.batches.create(\n",
    "    input_file_id=batch_input_file_id,\n",
    "    endpoint=\"/v1/chat/completions\",\n",
    "    completion_window=\"24h\",\n",
    "    metadata={\n",
    "      \"description\": filename.replace('.jsonl','')  # NOTE: it uses the file name as description\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50de36f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cancelling a batch\n",
    "# client.batches.cancel(\"batch_abc123\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "271d65ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Getting a list of all batches\n",
    "client.batches.list(limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "772191cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the status of a batch\n",
    "client.batches.retrieve(\"batch_676c3b98c7a0819085cf8a89ed7fdc67\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8a22c71",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieving the batch results:\n",
    "file_response = client.files.content(\"file-EqjRGjeDK7GZJ236nDbgdq\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2fe60e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert the output to randIndex_answer dictionary\n",
    "randIndex_answer = {}\n",
    "for line in file_response.text.splitlines():\n",
    "    item = json.loads(line)\n",
    "    randIndex = int(item['custom_id'])\n",
    "    answer = item['response']['body']['choices'][0]['message']['content']\n",
    "    randIndex_answer[randIndex] = answer\n",
    "print(len(randIndex_answer))\n",
    "print(randIndex_answer[4130])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b1e4c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the OpenAI answers\n",
    "save_to_json(randIndex_answer, f'answers_test_{PROMPT_TEMPLATE}.json')    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1325439c",
   "metadata": {},
   "source": [
    "# Use Open LLMs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64327e3b",
   "metadata": {},
   "source": [
    "This part is to reproduce the generated outputs by open LLMs (You can pass to the next section if you don't want to reproduce them). Before running the next cells in this section, you need to run the model in the vLLM using one of the following commands:\n",
    "\n",
    "```\n",
    "# you may need to run \"unset VLLM_ATTENTION_BACKEND\"\n",
    "python -m vllm.entrypoints.openai.api_server \\\n",
    "    --model meta-llama/Meta-Llama-3.1-8B-Instruct \\\n",
    "    --port 8000  \\\n",
    "    --gpu-memory-utilization 0.95 \\\n",
    "    --max-model-len 4096 \\\n",
    "    --download-dir /tmp/vllm_models \\\n",
    "    --api-key 123\n",
    "\n",
    "# you may need to run \"export VLLM_ATTENTION_BACKEND=FLASHINFER\" first.\n",
    "python -m vllm.entrypoints.openai.api_server \\\n",
    "    --model google/gemma-2-9b-it \\\n",
    "    --port 8000  \\\n",
    "    --gpu-memory-utilization 0.90 \\\n",
    "    --max-model-len 4096 \\\n",
    "    --download-dir /tmp/vllm_models \\\n",
    "    --api-key 123\n",
    "\n",
    "python -m vllm.entrypoints.openai.api_server \\\n",
    "    --model deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct \\\n",
    "    --port 8000  \\\n",
    "    --gpu-memory-utilization 0.90 \\\n",
    "    --max-model-len 4096 \\\n",
    "    --download-dir /tmp/vllm_models \\\n",
    "    --trust-remote-code \\\n",
    "    --api-key 123\n",
    "\n",
    "python -m vllm.entrypoints.openai.api_server \\\n",
    "    --model meta-llama/Meta-Llama-3.1-70B-Instruct \\\n",
    "    --port 8000  \\\n",
    "    --gpu-memory-utilization 0.90 \\\n",
    "    --max-model-len 4096 \\\n",
    "    --download-dir /tmp/vllm_models \\\n",
    "    --tensor-parallel-size 4 \\\n",
    "    --api-key 123\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c88689b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"123\"\n",
    "\n",
    "MODEL = 'Llama3.1-8B'\n",
    "MODEL_FullName = 'meta-llama/Meta-Llama-3.1-8B-Instruct'\n",
    "\n",
    "# MODEL = 'Llama3.1-70B'\n",
    "# MODEL_FullName = 'meta-llama/Meta-Llama-3.1-70B-Instruct'\n",
    "\n",
    "# MODEL = 'Gemma-2-9B'\n",
    "# MODEL_FullName = 'google/gemma-2-9b-it'\n",
    "\n",
    "# MODEL = 'DeepSeek-Coder-V2-Lite-Instruct'\n",
    "# MODEL_FullName = 'deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f8d7c34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test vllm to make sure it works\n",
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI(\n",
    "    api_key=\"123\",\n",
    "    base_url=\"http://localhost:8000/v1\",\n",
    ")\n",
    "\n",
    "prompts = ['300+30=','why is the sky blue?']\n",
    "\n",
    "# correct the format of prompts\n",
    "if 'llama3' in MODEL.lower():\n",
    "    prompts = [f\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n\\nYou are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n\" for prompt in prompts]\n",
    "elif 'deepseek' in MODEL.lower():\n",
    "    prompts = [f\"<｜begin▁of▁sentence｜>User: {prompt}\\n\\nAssistant:\" for prompt in prompts]\n",
    "elif 'openchat' in MODEL.lower() or 'starling' in MODEL.lower():\n",
    "    prompts = [f\"<s>GPT4 Correct User: {prompt}<|end_of_turn|>GPT4 Correct Assistant:\" for prompt in prompts]\n",
    "elif 'gemma' in MODEL.lower():\n",
    "    prompts = [f\"<bos><start_of_turn>user\\n{prompt}<end_of_turn>\\n<start_of_turn>model\\n\" for prompt in prompts]\n",
    "elif 'qwen' in MODEL.lower():\n",
    "    prompts = [f\"<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n<|im_start|>user\\n{prompt}<|im_end|>\\n<|im_start|>assistant\\n\" for prompt in prompts]\n",
    "elif 'mistral-nemo' in MODEL.lower():\n",
    "    prompts = [f\"<s>[INST]{prompt}[/INST]\" for prompt in prompts]\n",
    "#     prompts = [f\"<s>[INST]You are a helpful assistant.\\n\\n{prompt}[/INST]\" for prompt in prompts]\n",
    "elif 'phi' in MODEL.lower():\n",
    "    prompts = [f\"<|user|>\\n{prompt}<|end|>\\n<|assistant|>\\n\" for prompt in prompts]\n",
    "#     prompts = [f\"<|system|>\\nYou are a helpful assistant.<|end|>\\n<|user|>\\n{prompt}<|end|>\\n<|assistant|>\\n\" for prompt in prompts]\n",
    "else:\n",
    "    raise ValueError('Model not supported')    \n",
    "\n",
    "completion = client.completions.create(model=MODEL_FullName, prompt=prompts, max_tokens=2048, temperature=0.0, stop=[\"<|eot_id|>\"]) # we may need to set temperature to 0.01 rather 0.0\n",
    "# completion = client.completions.create(model=MODEL_FullName, prompt=prompts, max_tokens=2048, top_p=1, stop=[\"<|eot_id|>\"]) # it doesn't generate deterministic output\n",
    "\n",
    "for text in completion.choices:\n",
    "    print(text.text)\n",
    "    print('------------------------')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e212784",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_DIR = 'generated_outputs_new'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b3210f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the open model on the dataset. It splits the data to batches of size 100 and send them to the running model in vLLM.\n",
    "\n",
    "from openai import OpenAI\n",
    "import time\n",
    "\n",
    "client = OpenAI(\n",
    "    api_key=\"123\",\n",
    "    base_url=\"http://localhost:8000/v1\",\n",
    ")\n",
    "\n",
    "batch_size = 100\n",
    "batch_prompts = []\n",
    "batch_randIndex = []\n",
    "\n",
    "for i in range(0, len(df), batch_size):\n",
    "    batch_prompts.append(df.iloc[i:i + batch_size]['prompt2'].tolist())\n",
    "    batch_randIndex.append(df.iloc[i:i + batch_size]['rand_index'].tolist())\n",
    "print('Total number of batches:', len(batch_prompts))\n",
    "\n",
    "if 'openchat' in MODEL.lower() or 'starling' in MODEL.lower():\n",
    "    batch_prompts = [[f\"<s>GPT4 Correct User: {prompt}<|end_of_turn|>GPT4 Correct Assistant:\" for prompt in prompts] for prompts in batch_prompts]\n",
    "elif 'deepseek' in MODEL.lower():\n",
    "    batch_prompts = [[f\"<｜begin▁of▁sentence｜>User: {prompt}\\n\\nAssistant:\" for prompt in prompts] for prompts in batch_prompts]\n",
    "elif 'llama3' in MODEL.lower():\n",
    "    batch_prompts = [[f\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n\\nYou are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n\" for prompt in prompts] for prompts in batch_prompts]\n",
    "elif 'gemma' in MODEL.lower():\n",
    "    batch_prompts = [[f\"<bos><start_of_turn>user\\n{prompt}<end_of_turn>\\n<start_of_turn>model\\n\" for prompt in prompts] for prompts in batch_prompts]\n",
    "elif 'qwen' in MODEL.lower():\n",
    "    batch_prompts = [[f\"<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n<|im_start|>user\\n{prompt}<|im_end|>\\n<|im_start|>assistant\\n\" for prompt in prompts] for prompts in batch_prompts]\n",
    "elif 'mistral-nemo' in MODEL.lower():\n",
    "    batch_prompts = [[f\"<s>[INST]{prompt}[/INST]\" for prompt in prompts] for prompts in batch_prompts]\n",
    "    # batch_prompts = [[f\"<s>[INST]You are a helpful assistant.\\n\\n{prompt}[/INST]\" for prompt in prompts] for prompts in batch_prompts]\n",
    "elif 'phi' in MODEL.lower():\n",
    "    prompts = [f\"<|user|>\\n{prompt}<|end|>\\n<|assistant|>\\n\" for prompt in prompts]\n",
    "#     prompts = [f\"<|system|>\\nYou are a helpful assistant.<|end|>\\n<|user|>\\n{prompt}<|end|>\\n<|assistant|>\\n\" for prompt in prompts]\n",
    "else:\n",
    "    raise ValueError('Model not supported')\n",
    "\n",
    "start_time=time.time()\n",
    "randIndex_answer = {}\n",
    "unsaved_randIndex_answer = {}\n",
    "save_size = 10000\n",
    "for randIndex,prompts in zip(batch_randIndex,batch_prompts):\n",
    "    try:\n",
    "        print('Getting answers for the items',randIndex[0],'to',randIndex[-1]) # NOTE: set temperature to 0.0\n",
    "        try:\n",
    "            completion = client.completions.create(model=MODEL_FullName, prompt=prompts, max_tokens=2048, temperature=0.0, stop=[\"<|eot_id|>\"])\n",
    "        except Exception as e:\n",
    "            print('We failed with max_tokens=2048. Try with max_tokens=1024 ...')\n",
    "            completion = client.completions.create(model=MODEL_FullName, prompt=prompts, max_tokens=1024, temperature=0.0, stop=[\"<|eot_id|>\"])\n",
    "        for rindx,text in zip(randIndex,completion.choices):\n",
    "            randIndex_answer[rindx] = text.text\n",
    "            unsaved_randIndex_answer[rindx] = text.text\n",
    "        if len(unsaved_randIndex_answer)>=save_size:\n",
    "            if 'test' in DATASET:\n",
    "                filepath = OUTPUT_DIR + '/' + DATASET + '/' + MODEL + '/answers_test_' + PROMPT_TEMPLATE + '.json'\n",
    "            else:\n",
    "                filepath = OUTPUT_DIR + '/' + DATASET + '/' + MODEL + '/answers_'+str(list(unsaved_randIndex_answer.keys())[0]) + 'to' + str(list(unsaved_randIndex_answer.keys())[-1]) + '_' + PROMPT_TEMPLATE + '.json'\n",
    "            print('Saving items in', filepath)\n",
    "            save_to_json(unsaved_randIndex_answer, filepath)\n",
    "            unsaved_randIndex_answer = {}\n",
    "    except Exception as e:\n",
    "        print(\"An error occurred:\", e)\n",
    "\n",
    "# save the last part, if it is not the toy df\n",
    "if len(df)>500 and len(unsaved_randIndex_answer)>0:\n",
    "    if 'test' in DATASET:\n",
    "        filepath = OUTPUT_DIR + '/' + DATASET + '/' + MODEL + '/answers_test_' + PROMPT_TEMPLATE + '.json'\n",
    "    else:\n",
    "        filepath = OUTPUT_DIR + '/' + DATASET + '/' + MODEL + '/answers_'+str(list(unsaved_randIndex_answer.keys())[0]) + 'to' + str(list(unsaved_randIndex_answer.keys())[-1]) + '_' + PROMPT_TEMPLATE + '.json'\n",
    "    print('Saving items in', filepath)\n",
    "    save_to_json(unsaved_randIndex_answer, filepath)\n",
    "    unsaved_randIndex_answer = {}\n",
    "print('elapsed time:', int(time.time()-start_time), 'seconds')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4097f74d",
   "metadata": {},
   "source": [
    "# Determine the models and prompts for each dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3853cd8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determine the models and prompts for each dataset to find the corresponding filename in the output directory\n",
    "\n",
    "dataset_testpart = {'Python_test':'test', 'Java_test':'test', 'Mastropaolo':'test', 'Python':'10to2157', 'Java':'0to2003'}\n",
    "\n",
    "dataset_models = {'Mastropaolo':['gpt-4o-mini-2024-07-18', 'codeT5p-770m']} # RQ3\n",
    "for dataset in ['Java','Python']: # RQ2\n",
    "    dataset_models[dataset] = ['Llama3.1-8B', 'Gemma-2-9B', 'DeepSeek-Coder-V2-Lite-Instruct', 'Llama3.1-70B', 'gpt-4o-mini-2024-07-18']\n",
    "for dataset in ['Java_test','Python_test']: # RQ4\n",
    "    dataset_models[dataset] = ['Llama3.1-8B', 'Gemma-2-9B', 'DeepSeek-Coder-V2-Lite-Instruct', 'Llama3.1-70B', 'gpt-4o-mini-2024-07-18', 'codeT5p-220m-py', 'codeT5p-770m-py', 'codeT5p-220m', 'codeT5p-770m']\n",
    "    \n",
    "dataset_prompts = {'Mastropaolo':['NoExplain','TrainByControlled_TestByFiltered','TrainByFiltered_TestByFiltered']} # RQ3\n",
    "for dataset in ['Java','Python']: # RQ2\n",
    "    dataset_prompts[dataset] = ['Mastropaolo-T2', 'NoExplain', 'CoT1', 'CoT2']\n",
    "for dataset in ['Java_test','Python_test']: # RQ4\n",
    "    dataset_prompts[dataset] = ['Mastropaolo-T2', 'NoExplain', 'CoT1', 'CoT2', 'allTrain_lr1e-5', 'allTrain_lr5e-6'] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efdafb1b",
   "metadata": {},
   "source": [
    "# Check output and extract code from the generated answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8d808ae3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 3, 4, 8, 11, 13, 15, 16, 17, 22]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get a few rand_index for manual checking\n",
    "df.head(10)['rand_index'].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6aec7b8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "# load answers for check\n",
    "if RQ=='RQ2' or RQ=='RQ4':\n",
    "    randIndex_answer = load_from_json(f'generated_outputs/{DATASET}/DeepSeek-Coder-V2-Lite-Instruct/answers_{dataset_testpart[DATASET]}_CoT1.json')\n",
    "    # Convert keys to integers\n",
    "    randIndex_answer = {int(key): value for key, value in randIndex_answer.items()}\n",
    "    print(len(randIndex_answer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4aa17699",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "www.github.com/franzinc/agraph-java-client/commit/3580ee15ac8ddb5c935642314525162de55b5592\n",
      "src/com/franz/agraph/jena/AGInfGraph.java\n",
      "line: 102\n",
      "\n",
      " --------------- containing_method_before_repayment --------------- \n",
      "\n",
      "public ValidityReport validate() {\n",
      "\t// TODO Auto-generated method stub\n",
      "\treturn null;\n",
      "}\n",
      "\n",
      " --------------- containing_method_after_repayment --------------- \n",
      "\n",
      "public ValidityReport validate() {\n",
      "\tthrow new UnsupportedOperationException(AGUnsupportedOperation.message);\n",
      "}\n",
      "\n",
      " --------------- generated answer --------------- \n",
      "\n",
      " The Self-Applicable Technical Debt (SATD) comment \"TODO Auto-generated method stub\" indicates that the code contains a placeholder for a method implementation that should be replaced with actual logic. This is a common issue in codebases where developers might leave placeholders for methods they haven't yet implemented.\n",
      "\n",
      "To resolve this SATD, you should replace the placeholder with the actual logic that performs the validation. Here’s how you can update the code:\n",
      "\n",
      "### Updated Code:\n",
      "```java\n",
      "public ValidityReport validate() {\n",
      "    // Implement the validation logic here\n",
      "    // For example, you might want to check if certain conditions are met\n",
      "    // and return a ValidityReport based on those conditions\n",
      "    return new ValidityReport(); // Placeholder for actual implementation\n",
      "}\n",
      "```\n",
      "\n",
      "### Explanation:\n",
      "1. **Remove the TODO Comment**: The TODO comment is removed since the code now includes the actual logic.\n",
      "2. **Implement the Validation Logic**: Replace the placeholder return statement (`return new ValidityReport();`) with the actual logic that performs the validation. This might involve checking the state of the object, fetching data from a database, or performing other necessary operations to determine the validity of the object.\n",
      "\n",
      "By implementing the validation logic, you are resolving the SATD comment and making the code more functional and useful.\n"
     ]
    }
   ],
   "source": [
    "# check one answer\n",
    "indx = 22\n",
    "print(df[df['rand_index']==indx]['deletion_commit_url'].tolist()[0])\n",
    "print(df[df['rand_index']==indx]['last_appeared_in_file'].tolist()[0])\n",
    "print('line:', df[df['rand_index']==indx]['last_appeared_in_line'].tolist()[0])\n",
    "print('\\n', '-'*15, 'containing_method_before_repayment', '-'*15, '\\n')\n",
    "print(df[df['rand_index']==indx]['containing_method_before_repayment'].tolist()[0])\n",
    "print('\\n', '-'*15, 'containing_method_after_repayment', '-'*15, '\\n')\n",
    "print(df[df['rand_index']==indx]['containing_method_after_repayment'].tolist()[0])\n",
    "print('\\n', '-'*15, 'generated answer', '-'*15, '\\n')\n",
    "print(randIndex_answer[indx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "af669d6b",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      " To resolve the SATD (Sticky Assignee Task Description), you need to replace the `TODO` comment with the actual implementation of the `readConnectionContent(connectionURI)` method. This method is currently a placeholder (`return null;`).\n",
      "\n",
      "Here's the updated code:\n",
      "\n",
      "### Updated Code:\n",
      "```java\n",
      "public Connection readConnection(URI connectionURI) throws NoSuchConnectionException {\n",
      "    logger.debug(MessageFormat.format(\"need-facing: READ_CONNECTION called for connection {0}\", connectionURI));\n",
      "\n",
      "    // Implement readConnectionContent(connectionURI)\n",
      "    return readConnectionContent(connectionURI);\n",
      "}\n",
      "\n",
      "// Assuming readConnectionContent(URI connectionURI) is defined elsewhere in the class\n",
      "private Connection readConnectionContent(URI connectionURI) throws NoSuchConnectionException {\n",
      "    // Implementation of reading the connection content\n",
      "    // This is a placeholder implementation\n",
      "    // You need to replace this with the actual logic\n",
      "    throw new NoSuchConnectionException(\"Not implemented\");\n",
      "}\n",
      "```\n",
      "\n",
      "### Explanation:\n",
      "1. **Replace the TODO comment**: The `TODO` comment in the original code indicates that `readConnectionContent(connectionURI)` needs to be implemented. In the updated code, I have added a placeholder implementation of `readConnectionContent(connectionURI)` within the same class.\n",
      "2. **Implement the method**: The `readConnectionContent(URI connectionURI)` method is a placeholder that throws a `NoSuchConnectionException`. You need to replace this with the actual logic to read the connection content based on the `connectionURI`.\n",
      "\n",
      "This approach ensures that the `TODO` comment is resolved and the method is properly implemented.\n",
      "------------------------------\n",
      "public Connection readConnection(URI connectionURI) throws NoSuchConnectionException {\n",
      "    logger.debug(MessageFormat.format(\"need-facing: READ_CONNECTION called for connection {0}\", connectionURI));\n",
      "\n",
      "    // Implement readConnectionContent(connectionURI)\n",
      "    return readConnectionContent(connectionURI);\n",
      "}\n",
      "\n",
      "// Assuming readConnectionContent(URI connectionURI) is defined elsewhere in the class\n",
      "private Connection readConnectionContent(URI connectionURI) throws NoSuchConnectionException {\n",
      "    // Implementation of reading the connection content\n",
      "    // This is a placeholder implementation\n",
      "    // You need to replace this with the actual logic\n",
      "    throw new NoSuchConnectionException(\"Not implemented\");\n",
      "}\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# check extracting code from answers\n",
    "for i in df.head(1)['rand_index']:\n",
    "    print(i)\n",
    "    print(randIndex_answer[i])\n",
    "    print('-'*30)\n",
    "    print(extract_code(randIndex_answer[i], LANGUAGE))\n",
    "    print('='*50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40f14b33",
   "metadata": {},
   "source": [
    "# Check how remove ICD (imports, comments, and docstings) works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5a108a97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def example_function():\n",
      "    x = 10\n",
      "    y = 20\n",
      "    a = '#Hello#'\n",
      "    print('###Hello###')\n",
      "    return x + y\n"
     ]
    }
   ],
   "source": [
    "# Example usage for remove_docstrings_and_comments_from_python\n",
    "\n",
    "python_code = \"\"\"\n",
    "import os\n",
    "from os import xyz\n",
    "# This is a sample Python code\n",
    "#\n",
    "def example_function():\n",
    "    '''This is docstring'''\n",
    "    x = 10  # Initialize x with 10\n",
    "    y = 20  # Initialize y with 20\n",
    "    a = \"#Hello#\"  # Initialize a with \"#Hello#\"\n",
    "    print(\"###Hello###\")\n",
    "    # The following line does the sum\n",
    "    return x + y  # Return the sum\n",
    "\"\"\"\n",
    "\n",
    "clean_code = remove_docstrings_and_comments_from_python(remove_imports(python_code))\n",
    "print(clean_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b539376e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "public class Example {\n",
      "    public void method() {\n",
      "        int x = 1; \n",
      "        int y = 2; \n",
      "    }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Example usage for remove_comments_and_javadoc_from_java\n",
    "java_code = \"\"\"\n",
    "public class Example {\n",
    "    /**\n",
    "     * This is a Javadoc comment\n",
    "     */\n",
    "    public void method() {\n",
    "        // This is a single-line comment\n",
    "        int x = 1; /* This is a multi-line comment */\n",
    "        int y = 2; // This is another // single-line comment\n",
    "    }\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "cleaned_code = remove_comments_and_javadoc_from_java(java_code)\n",
    "print(cleaned_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6432f0a",
   "metadata": {},
   "source": [
    "# Get Exact Match Scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2356c03e",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_DIR = 'generated_outputs'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "26e31161",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Llama3.1-8B with Mastropaolo-T2\n",
      "  Exact Match without any changes: 20\n",
      "  Exact Match by ignoring whitespaces: 30\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 36\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 33\n",
      "      Items with three or more added lines in ground truth: 3\n",
      "Average of Levenshtein Distance: 198.639\n",
      "--------------------------------------------------\n",
      "Llama3.1-8B with NoExplain\n",
      "  Exact Match without any changes: 28\n",
      "  Exact Match by ignoring whitespaces: 39\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 57\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 54\n",
      "      Items with three or more added lines in ground truth: 3\n",
      "Average of Levenshtein Distance: 220.384\n",
      "--------------------------------------------------\n",
      "Llama3.1-8B with CoT1\n",
      "  Exact Match without any changes: 22\n",
      "  Exact Match by ignoring whitespaces: 30\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 54\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 51\n",
      "      Items with three or more added lines in ground truth: 3\n",
      "Average of Levenshtein Distance: 283.399\n",
      "--------------------------------------------------\n",
      "Llama3.1-8B with CoT2\n",
      "  Exact Match without any changes: 21\n",
      "  Exact Match by ignoring whitespaces: 29\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 53\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 48\n",
      "      Items with three or more added lines in ground truth: 5\n",
      "Average of Levenshtein Distance: 292.799\n",
      "--------------------------------------------------\n",
      "==================================================\n",
      "Gemma-2-9B with Mastropaolo-T2\n",
      "  Exact Match without any changes: 29\n",
      "  Exact Match by ignoring whitespaces: 45\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 58\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 55\n",
      "      Items with three or more added lines in ground truth: 3\n",
      "Average of Levenshtein Distance: 190.973\n",
      "--------------------------------------------------\n",
      "Gemma-2-9B with NoExplain\n",
      "  Exact Match without any changes: 47\n",
      "  Exact Match by ignoring whitespaces: 70\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 88\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 83\n",
      "      Items with three or more added lines in ground truth: 5\n",
      "Average of Levenshtein Distance: 209.015\n",
      "--------------------------------------------------\n",
      "Gemma-2-9B with CoT1\n",
      "  Exact Match without any changes: 11\n",
      "  Exact Match by ignoring whitespaces: 22\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 58\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 53\n",
      "      Items with three or more added lines in ground truth: 5\n",
      "Average of Levenshtein Distance: 335.632\n",
      "--------------------------------------------------\n",
      "Gemma-2-9B with CoT2\n",
      "  Exact Match without any changes: 16\n",
      "  Exact Match by ignoring whitespaces: 24\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 58\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 53\n",
      "      Items with three or more added lines in ground truth: 5\n",
      "Average of Levenshtein Distance: 318.087\n",
      "--------------------------------------------------\n",
      "==================================================\n",
      "DeepSeek-Coder-V2-Lite-Instruct with Mastropaolo-T2\n",
      "  Exact Match without any changes: 22\n",
      "  Exact Match by ignoring whitespaces: 31\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 37\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 34\n",
      "      Items with three or more added lines in ground truth: 3\n",
      "Average of Levenshtein Distance: 197.232\n",
      "--------------------------------------------------\n",
      "DeepSeek-Coder-V2-Lite-Instruct with NoExplain\n",
      "  Exact Match without any changes: 32\n",
      "  Exact Match by ignoring whitespaces: 39\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 62\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 56\n",
      "      Items with three or more added lines in ground truth: 6\n",
      "Average of Levenshtein Distance: 205.433\n",
      "--------------------------------------------------\n",
      "DeepSeek-Coder-V2-Lite-Instruct with CoT1\n",
      "  Exact Match without any changes: 22\n",
      "  Exact Match by ignoring whitespaces: 26\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 57\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 55\n",
      "      Items with three or more added lines in ground truth: 2\n",
      "Average of Levenshtein Distance: 239.629\n",
      "--------------------------------------------------\n",
      "DeepSeek-Coder-V2-Lite-Instruct with CoT2\n",
      "  Exact Match without any changes: 15\n",
      "  Exact Match by ignoring whitespaces: 22\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 54\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 52\n",
      "      Items with three or more added lines in ground truth: 2\n",
      "Average of Levenshtein Distance: 259.156\n",
      "--------------------------------------------------\n",
      "==================================================\n",
      "Llama3.1-70B with Mastropaolo-T2\n",
      "  Exact Match without any changes: 35\n",
      "  Exact Match by ignoring whitespaces: 55\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 71\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 66\n",
      "      Items with three or more added lines in ground truth: 5\n",
      "Average of Levenshtein Distance: 259.126\n",
      "--------------------------------------------------\n",
      "Llama3.1-70B with NoExplain\n",
      "  Exact Match without any changes: 43\n",
      "  Exact Match by ignoring whitespaces: 62\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 74\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 68\n",
      "      Items with three or more added lines in ground truth: 6\n",
      "Average of Levenshtein Distance: 228.128\n",
      "--------------------------------------------------\n",
      "Llama3.1-70B with CoT1\n",
      "  Exact Match without any changes: 32\n",
      "  Exact Match by ignoring whitespaces: 50\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 76\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 70\n",
      "      Items with three or more added lines in ground truth: 6\n",
      "Average of Levenshtein Distance: 330.547\n",
      "--------------------------------------------------\n",
      "Llama3.1-70B with CoT2\n",
      "  Exact Match without any changes: 29\n",
      "  Exact Match by ignoring whitespaces: 46\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 72\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 64\n",
      "      Items with three or more added lines in ground truth: 8\n",
      "Average of Levenshtein Distance: 323.855\n",
      "--------------------------------------------------\n",
      "==================================================\n",
      "gpt-4o-mini-2024-07-18 with Mastropaolo-T2\n",
      "  Exact Match without any changes: 10\n",
      "  Exact Match by ignoring whitespaces: 13\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 29\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 28\n",
      "      Items with three or more added lines in ground truth: 1\n",
      "Average of Levenshtein Distance: 210.557\n",
      "--------------------------------------------------\n",
      "gpt-4o-mini-2024-07-18 with NoExplain\n",
      "  Exact Match without any changes: 36\n",
      "  Exact Match by ignoring whitespaces: 49\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 87\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 78\n",
      "      Items with three or more added lines in ground truth: 9\n",
      "Average of Levenshtein Distance: 197.862\n",
      "--------------------------------------------------\n",
      "gpt-4o-mini-2024-07-18 with CoT1\n",
      "  Exact Match without any changes: 5\n",
      "  Exact Match by ignoring whitespaces: 8\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 74\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 66\n",
      "      Items with three or more added lines in ground truth: 8\n",
      "Average of Levenshtein Distance: 242.811\n",
      "--------------------------------------------------\n",
      "gpt-4o-mini-2024-07-18 with CoT2\n",
      "  Exact Match without any changes: 6\n",
      "  Exact Match by ignoring whitespaces: 9\n",
      "  Ignoring imports, code comments, and docstrings:\n",
      "    Exact Match (all): 72\n",
      "    Difficulty group:\n",
      "      Items with at most two added lines in ground truth: 64\n",
      "      Items with three or more added lines in ground truth: 8\n",
      "Average of Levenshtein Distance: 259.421\n",
      "--------------------------------------------------\n",
      "==================================================\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>Model</th>\n",
       "      <th>Template</th>\n",
       "      <th>Avg deletes</th>\n",
       "      <th>Avg inserts</th>\n",
       "      <th>EM</th>\n",
       "      <th>EM on Easy</th>\n",
       "      <th>EM on Hard</th>\n",
       "      <th>Avg Levenshtein Distance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>2.161</td>\n",
       "      <td>2.461</td>\n",
       "      <td>0.036</td>\n",
       "      <td>0.033</td>\n",
       "      <td>0.003</td>\n",
       "      <td>198.639</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>4.391</td>\n",
       "      <td>3.538</td>\n",
       "      <td>0.057</td>\n",
       "      <td>0.054</td>\n",
       "      <td>0.003</td>\n",
       "      <td>220.384</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>7.317</td>\n",
       "      <td>5.468</td>\n",
       "      <td>0.054</td>\n",
       "      <td>0.051</td>\n",
       "      <td>0.003</td>\n",
       "      <td>283.399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>8.682</td>\n",
       "      <td>5.091</td>\n",
       "      <td>0.053</td>\n",
       "      <td>0.048</td>\n",
       "      <td>0.005</td>\n",
       "      <td>292.799</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>1.795</td>\n",
       "      <td>2.606</td>\n",
       "      <td>0.058</td>\n",
       "      <td>0.055</td>\n",
       "      <td>0.003</td>\n",
       "      <td>190.973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>2.718</td>\n",
       "      <td>4.145</td>\n",
       "      <td>0.088</td>\n",
       "      <td>0.083</td>\n",
       "      <td>0.005</td>\n",
       "      <td>209.015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>7.638</td>\n",
       "      <td>8.328</td>\n",
       "      <td>0.058</td>\n",
       "      <td>0.053</td>\n",
       "      <td>0.005</td>\n",
       "      <td>335.632</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>7.398</td>\n",
       "      <td>7.895</td>\n",
       "      <td>0.058</td>\n",
       "      <td>0.053</td>\n",
       "      <td>0.005</td>\n",
       "      <td>318.087</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>1.322</td>\n",
       "      <td>2.455</td>\n",
       "      <td>0.037</td>\n",
       "      <td>0.034</td>\n",
       "      <td>0.003</td>\n",
       "      <td>197.232</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>4.009</td>\n",
       "      <td>3.697</td>\n",
       "      <td>0.062</td>\n",
       "      <td>0.056</td>\n",
       "      <td>0.006</td>\n",
       "      <td>205.433</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>7.576</td>\n",
       "      <td>4.352</td>\n",
       "      <td>0.057</td>\n",
       "      <td>0.055</td>\n",
       "      <td>0.002</td>\n",
       "      <td>239.629</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>9.372</td>\n",
       "      <td>4.634</td>\n",
       "      <td>0.054</td>\n",
       "      <td>0.052</td>\n",
       "      <td>0.002</td>\n",
       "      <td>259.156</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>5.428</td>\n",
       "      <td>5.123</td>\n",
       "      <td>0.071</td>\n",
       "      <td>0.066</td>\n",
       "      <td>0.005</td>\n",
       "      <td>259.126</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>5.624</td>\n",
       "      <td>4.459</td>\n",
       "      <td>0.074</td>\n",
       "      <td>0.068</td>\n",
       "      <td>0.006</td>\n",
       "      <td>228.128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>8.231</td>\n",
       "      <td>7.622</td>\n",
       "      <td>0.076</td>\n",
       "      <td>0.070</td>\n",
       "      <td>0.006</td>\n",
       "      <td>330.547</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>8.313</td>\n",
       "      <td>7.474</td>\n",
       "      <td>0.072</td>\n",
       "      <td>0.064</td>\n",
       "      <td>0.008</td>\n",
       "      <td>323.855</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>3.181</td>\n",
       "      <td>2.477</td>\n",
       "      <td>0.029</td>\n",
       "      <td>0.028</td>\n",
       "      <td>0.001</td>\n",
       "      <td>210.557</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>4.535</td>\n",
       "      <td>3.779</td>\n",
       "      <td>0.087</td>\n",
       "      <td>0.078</td>\n",
       "      <td>0.009</td>\n",
       "      <td>197.862</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>9.356</td>\n",
       "      <td>4.519</td>\n",
       "      <td>0.074</td>\n",
       "      <td>0.066</td>\n",
       "      <td>0.008</td>\n",
       "      <td>242.811</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>10.125</td>\n",
       "      <td>4.713</td>\n",
       "      <td>0.072</td>\n",
       "      <td>0.064</td>\n",
       "      <td>0.008</td>\n",
       "      <td>259.421</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              Model        Template  Avg deletes  Avg inserts  \\\n",
       "0                       Llama3.1-8B  Mastropaolo-T2        2.161        2.461   \n",
       "1                       Llama3.1-8B       NoExplain        4.391        3.538   \n",
       "2                       Llama3.1-8B            CoT1        7.317        5.468   \n",
       "3                       Llama3.1-8B            CoT2        8.682        5.091   \n",
       "4                        Gemma-2-9B  Mastropaolo-T2        1.795        2.606   \n",
       "5                        Gemma-2-9B       NoExplain        2.718        4.145   \n",
       "6                        Gemma-2-9B            CoT1        7.638        8.328   \n",
       "7                        Gemma-2-9B            CoT2        7.398        7.895   \n",
       "8   DeepSeek-Coder-V2-Lite-Instruct  Mastropaolo-T2        1.322        2.455   \n",
       "9   DeepSeek-Coder-V2-Lite-Instruct       NoExplain        4.009        3.697   \n",
       "10  DeepSeek-Coder-V2-Lite-Instruct            CoT1        7.576        4.352   \n",
       "11  DeepSeek-Coder-V2-Lite-Instruct            CoT2        9.372        4.634   \n",
       "12                     Llama3.1-70B  Mastropaolo-T2        5.428        5.123   \n",
       "13                     Llama3.1-70B       NoExplain        5.624        4.459   \n",
       "14                     Llama3.1-70B            CoT1        8.231        7.622   \n",
       "15                     Llama3.1-70B            CoT2        8.313        7.474   \n",
       "16           gpt-4o-mini-2024-07-18  Mastropaolo-T2        3.181        2.477   \n",
       "17           gpt-4o-mini-2024-07-18       NoExplain        4.535        3.779   \n",
       "18           gpt-4o-mini-2024-07-18            CoT1        9.356        4.519   \n",
       "19           gpt-4o-mini-2024-07-18            CoT2       10.125        4.713   \n",
       "\n",
       "       EM  EM on Easy  EM on Hard  Avg Levenshtein Distance  \n",
       "0   0.036       0.033       0.003                   198.639  \n",
       "1   0.057       0.054       0.003                   220.384  \n",
       "2   0.054       0.051       0.003                   283.399  \n",
       "3   0.053       0.048       0.005                   292.799  \n",
       "4   0.058       0.055       0.003                   190.973  \n",
       "5   0.088       0.083       0.005                   209.015  \n",
       "6   0.058       0.053       0.005                   335.632  \n",
       "7   0.058       0.053       0.005                   318.087  \n",
       "8   0.037       0.034       0.003                   197.232  \n",
       "9   0.062       0.056       0.006                   205.433  \n",
       "10  0.057       0.055       0.002                   239.629  \n",
       "11  0.054       0.052       0.002                   259.156  \n",
       "12  0.071       0.066       0.005                   259.126  \n",
       "13  0.074       0.068       0.006                   228.128  \n",
       "14  0.076       0.070       0.006                   330.547  \n",
       "15  0.072       0.064       0.008                   323.855  \n",
       "16  0.029       0.028       0.001                   210.557  \n",
       "17  0.087       0.078       0.009                   197.862  \n",
       "18  0.074       0.066       0.008                   242.811  \n",
       "19  0.072       0.064       0.008                   259.421  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_template_em = {}\n",
    "model_randIndex_answer = {}\n",
    "data = []\n",
    "for model_name in dataset_models[DATASET]: \n",
    "    model_template_em[model_name] = {}\n",
    "    for prompt_template in dataset_prompts[DATASET]:\n",
    "        file_path = f'{OUTPUT_DIR}/{DATASET}/{model_name}/answers_{dataset_testpart[DATASET]}_{prompt_template}.json'\n",
    "        if os.path.exists(file_path)==False:\n",
    "            continue\n",
    "        randIndex_answer = load_from_json(file_path)\n",
    "        randIndex_answer = {int(key): value for key, value in randIndex_answer.items()}\n",
    "        randIndex_answer = {k: v for k, v in randIndex_answer.items() if k in set(df['rand_index'])} # filter out items not exist in df\n",
    "        model_randIndex_answer[model_name] = randIndex_answer\n",
    "        if 'codet5' in model_name.lower():\n",
    "            extract_code_from_answer=False\n",
    "        else:\n",
    "            extract_code_from_answer=True            \n",
    "        print(model_name, 'with', prompt_template)\n",
    "        randIndex_distance, em, black_failed, black_failed_em = get_exact_matches(df, randIndex_answer, LANGUAGE, DATASET, extract_code_from_answer, ignore_whitespace=False, format_code=False)\n",
    "        print('  Exact Match without any changes:',len(em))\n",
    "        randIndex_distance, em, black_failed, black_failed_em = get_exact_matches(df, randIndex_answer, LANGUAGE, DATASET, extract_code_from_answer, ignore_whitespace=True)\n",
    "        print('  Exact Match by ignoring whitespaces:',len(em))\n",
    "        if False:\n",
    "            print('  Considering imports, code comments, and docstrings:')\n",
    "            randIndex_distance, em, black_failed, black_failed_em = get_exact_matches(df, randIndex_answer, LANGUAGE, DATASET, extract_code_from_answer)\n",
    "            print('    Exact Match (all):',len(em), list(em.keys()))\n",
    "            print('    Failed by black:',len(black_failed))\n",
    "            print('    Failed by black but is exact match by ignoring whitespaces:',len(black_failed_em), list(black_failed_em.keys()))\n",
    "        print('  Ignoring imports, code comments, and docstrings:')\n",
    "        randIndex_distance, em, black_failed, black_failed_em = get_exact_matches(df, randIndex_answer, LANGUAGE, DATASET, extract_code_from_answer, ignore_docstrings_and_comments=True)\n",
    "        model_template_em[model_name][prompt_template] = set(em.keys())\n",
    "        print('    Exact Match (all):',len(em)) # , list(em.keys()))\n",
    "        if LANGUAGE=='Python':\n",
    "            print('    Failed by black:',len(black_failed))\n",
    "            print('    Failed by black but is exact match by ignoring whitespaces:',len(black_failed_em), list(black_failed_em.keys()))\n",
    "        print('    Difficulty group:')\n",
    "        items_easy = get_itmes_having_specific_number_of_inserted_lines(df,em.keys(),0,2)\n",
    "        print('      Items with at most two added lines in ground truth:', len(items_easy))\n",
    "        items_hard = get_itmes_having_specific_number_of_inserted_lines(df,em.keys(),3,1000)\n",
    "        print('      Items with three or more added lines in ground truth:', len(items_hard))\n",
    "        avg_leven_dist = sum(randIndex_distance.values()) / len(randIndex_distance)\n",
    "        print('Average of Levenshtein Distance:', avg_leven_dist)\n",
    "        # get the average number of deleted and inserted lines\n",
    "        deletes = []\n",
    "        inserts = []\n",
    "        for indx, answer in randIndex_answer.items():\n",
    "            if extract_code_from_answer:\n",
    "                answercode = extract_code(answer,LANGUAGE)\n",
    "            else:\n",
    "                answercode = answer\n",
    "            inputcode = df[df['rand_index']==indx]['containing_method_before_repayment'].tolist()[0]\n",
    "            deleted,inserted = get_deleted_and_inserted_lines(answercode,inputcode)\n",
    "            deletes.append(len(deleted))\n",
    "            inserts.append(len(inserted))\n",
    "        avg_deletes = sum(deletes)/len(deletes)\n",
    "        avg_inserts = sum(inserts)/len(inserts)\n",
    "        print('-'*50)\n",
    "        data.append([model_name, prompt_template, avg_deletes, avg_inserts, len(em)/len(randIndex_answer), len(items_easy)/len(randIndex_answer), len(items_hard)/len(randIndex_answer), avg_leven_dist])\n",
    "    print('='*50)\n",
    "    \n",
    "# convert data to df\n",
    "df_EM = pd.DataFrame(data, columns=['Model', 'Template', 'Avg deletes', 'Avg inserts', 'EM', 'EM on Easy', 'EM on Hard', 'Avg Levenshtein Distance'])\n",
    "display(df_EM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "51422d2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Llama3.1-8B': 0.1, 'Gemma-2-9B': 0.106, 'DeepSeek-Coder-V2-Lite-Instruct': 0.093, 'Llama3.1-70B': 0.111, 'gpt-4o-mini-2024-07-18': 0.106}\n"
     ]
    }
   ],
   "source": [
    "# Optimal Performance with an Oracle Template (Discussion 6.1)\n",
    "model_oracle = {}\n",
    "if RQ=='RQ2':\n",
    "    for model_name in dataset_models[DATASET]:\n",
    "        model_oracle[model_name] = set()\n",
    "        for prompt_template in dataset_prompts[DATASET]:\n",
    "            model_oracle[model_name] = model_oracle[model_name] | model_template_em[model_name][prompt_template]\n",
    "    print({model:len(oracle)/len(df) for model,oracle in model_oracle.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b60828",
   "metadata": {},
   "source": [
    "# BLEU and CrystalBLEU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "436f75db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('=',): 3654\n",
      "('{',): 3653\n",
      "('}',): 3625\n",
      "('if',): 1452\n",
      "('new',): 1156\n",
      "('//',): 1011\n",
      "('return',): 922\n",
      "('public',): 775\n",
      "('}', '}'): 762\n",
      "('+',): 756\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "from nltk.util import ngrams\n",
    "\n",
    "k=500\n",
    "\n",
    "tokenized_corpus = [line.split() for line in df['containing_method_after_repayment'].tolist()]\n",
    "# convert list_of_lists to list\n",
    "tokenized_corpus = [item for sublist in tokenized_corpus for item in sublist]\n",
    "all_ngrams = []\n",
    "for n in range(1, 5):\n",
    "    all_ngrams.extend(list(ngrams(tokenized_corpus, n)))\n",
    "\n",
    "frequencies = Counter(all_ngrams)\n",
    "trivially_shared_ngrams = dict(frequencies.most_common(k))\n",
    "# show the first n itmes\n",
    "for k, v in list(trivially_shared_ngrams.items())[:10]: print(f\"{k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e153b465",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jovyan/conda-envs/SATD_repayment_py3.10/lib/python3.10/site-packages/crystalbleu.py:565: UserWarning: \n",
      "The hypothesis contains 0 counts of 2-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n",
      "/home/jovyan/conda-envs/SATD_repayment_py3.10/lib/python3.10/site-packages/crystalbleu.py:565: UserWarning: \n",
      "The hypothesis contains 0 counts of 3-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n",
      "/home/jovyan/conda-envs/SATD_repayment_py3.10/lib/python3.10/site-packages/crystalbleu.py:565: UserWarning: \n",
      "The hypothesis contains 0 counts of 4-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>Model</th>\n",
       "      <th>Template</th>\n",
       "      <th>BLEU</th>\n",
       "      <th>CrystalBLEU</th>\n",
       "      <th>LineP</th>\n",
       "      <th>lineR</th>\n",
       "      <th>lineF</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>0.571558</td>\n",
       "      <td>0.609780</td>\n",
       "      <td>0.685796</td>\n",
       "      <td>0.635833</td>\n",
       "      <td>0.636843</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>0.580184</td>\n",
       "      <td>0.634841</td>\n",
       "      <td>0.646905</td>\n",
       "      <td>0.637976</td>\n",
       "      <td>0.624444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>0.521046</td>\n",
       "      <td>0.578498</td>\n",
       "      <td>0.579323</td>\n",
       "      <td>0.600717</td>\n",
       "      <td>0.565996</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>0.514933</td>\n",
       "      <td>0.571604</td>\n",
       "      <td>0.549777</td>\n",
       "      <td>0.605577</td>\n",
       "      <td>0.553629</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>0.584565</td>\n",
       "      <td>0.632057</td>\n",
       "      <td>0.736707</td>\n",
       "      <td>0.655819</td>\n",
       "      <td>0.671136</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>0.583495</td>\n",
       "      <td>0.636123</td>\n",
       "      <td>0.710583</td>\n",
       "      <td>0.633125</td>\n",
       "      <td>0.645579</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>0.469336</td>\n",
       "      <td>0.527929</td>\n",
       "      <td>0.554741</td>\n",
       "      <td>0.544667</td>\n",
       "      <td>0.519439</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>0.482686</td>\n",
       "      <td>0.541188</td>\n",
       "      <td>0.560446</td>\n",
       "      <td>0.555850</td>\n",
       "      <td>0.529159</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>0.571782</td>\n",
       "      <td>0.603556</td>\n",
       "      <td>0.715122</td>\n",
       "      <td>0.639117</td>\n",
       "      <td>0.650576</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>0.582480</td>\n",
       "      <td>0.627090</td>\n",
       "      <td>0.662053</td>\n",
       "      <td>0.626410</td>\n",
       "      <td>0.621741</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>0.558710</td>\n",
       "      <td>0.600962</td>\n",
       "      <td>0.593795</td>\n",
       "      <td>0.612780</td>\n",
       "      <td>0.581403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>0.539122</td>\n",
       "      <td>0.583910</td>\n",
       "      <td>0.550109</td>\n",
       "      <td>0.606296</td>\n",
       "      <td>0.556113</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>0.542808</td>\n",
       "      <td>0.583036</td>\n",
       "      <td>0.640029</td>\n",
       "      <td>0.612951</td>\n",
       "      <td>0.600548</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>0.571903</td>\n",
       "      <td>0.620899</td>\n",
       "      <td>0.618740</td>\n",
       "      <td>0.629207</td>\n",
       "      <td>0.605174</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>0.495766</td>\n",
       "      <td>0.543411</td>\n",
       "      <td>0.560374</td>\n",
       "      <td>0.575503</td>\n",
       "      <td>0.539061</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>0.495997</td>\n",
       "      <td>0.544912</td>\n",
       "      <td>0.556674</td>\n",
       "      <td>0.579730</td>\n",
       "      <td>0.539682</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>0.564011</td>\n",
       "      <td>0.596847</td>\n",
       "      <td>0.682929</td>\n",
       "      <td>0.643040</td>\n",
       "      <td>0.638282</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>0.594803</td>\n",
       "      <td>0.641592</td>\n",
       "      <td>0.656432</td>\n",
       "      <td>0.650165</td>\n",
       "      <td>0.635587</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>0.549845</td>\n",
       "      <td>0.589923</td>\n",
       "      <td>0.553251</td>\n",
       "      <td>0.622921</td>\n",
       "      <td>0.565938</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>0.539992</td>\n",
       "      <td>0.580590</td>\n",
       "      <td>0.539862</td>\n",
       "      <td>0.616909</td>\n",
       "      <td>0.554579</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              Model        Template      BLEU  CrystalBLEU  \\\n",
       "0                       Llama3.1-8B  Mastropaolo-T2  0.571558     0.609780   \n",
       "1                       Llama3.1-8B       NoExplain  0.580184     0.634841   \n",
       "2                       Llama3.1-8B            CoT1  0.521046     0.578498   \n",
       "3                       Llama3.1-8B            CoT2  0.514933     0.571604   \n",
       "4                        Gemma-2-9B  Mastropaolo-T2  0.584565     0.632057   \n",
       "5                        Gemma-2-9B       NoExplain  0.583495     0.636123   \n",
       "6                        Gemma-2-9B            CoT1  0.469336     0.527929   \n",
       "7                        Gemma-2-9B            CoT2  0.482686     0.541188   \n",
       "8   DeepSeek-Coder-V2-Lite-Instruct  Mastropaolo-T2  0.571782     0.603556   \n",
       "9   DeepSeek-Coder-V2-Lite-Instruct       NoExplain  0.582480     0.627090   \n",
       "10  DeepSeek-Coder-V2-Lite-Instruct            CoT1  0.558710     0.600962   \n",
       "11  DeepSeek-Coder-V2-Lite-Instruct            CoT2  0.539122     0.583910   \n",
       "12                     Llama3.1-70B  Mastropaolo-T2  0.542808     0.583036   \n",
       "13                     Llama3.1-70B       NoExplain  0.571903     0.620899   \n",
       "14                     Llama3.1-70B            CoT1  0.495766     0.543411   \n",
       "15                     Llama3.1-70B            CoT2  0.495997     0.544912   \n",
       "16           gpt-4o-mini-2024-07-18  Mastropaolo-T2  0.564011     0.596847   \n",
       "17           gpt-4o-mini-2024-07-18       NoExplain  0.594803     0.641592   \n",
       "18           gpt-4o-mini-2024-07-18            CoT1  0.549845     0.589923   \n",
       "19           gpt-4o-mini-2024-07-18            CoT2  0.539992     0.580590   \n",
       "\n",
       "       LineP     lineR     lineF  \n",
       "0   0.685796  0.635833  0.636843  \n",
       "1   0.646905  0.637976  0.624444  \n",
       "2   0.579323  0.600717  0.565996  \n",
       "3   0.549777  0.605577  0.553629  \n",
       "4   0.736707  0.655819  0.671136  \n",
       "5   0.710583  0.633125  0.645579  \n",
       "6   0.554741  0.544667  0.519439  \n",
       "7   0.560446  0.555850  0.529159  \n",
       "8   0.715122  0.639117  0.650576  \n",
       "9   0.662053  0.626410  0.621741  \n",
       "10  0.593795  0.612780  0.581403  \n",
       "11  0.550109  0.606296  0.556113  \n",
       "12  0.640029  0.612951  0.600548  \n",
       "13  0.618740  0.629207  0.605174  \n",
       "14  0.560374  0.575503  0.539061  \n",
       "15  0.556674  0.579730  0.539682  \n",
       "16  0.682929  0.643040  0.638282  \n",
       "17  0.656432  0.650165  0.635587  \n",
       "18  0.553251  0.622921  0.565938  \n",
       "19  0.539862  0.616909  0.554579  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Baseline - Use containing_method_before_repayment as candidate:\n",
      "  Average BLEU: 0.564\n",
      "  Average crystal BLEU: 0.611\n",
      "  Average lineP: 0.740\n",
      "  Average lineR: 0.655\n",
      "  Average lineF: 0.669\n"
     ]
    }
   ],
   "source": [
    "# NOTE: run this cell only for RQ2\n",
    "# Calculate BLEU, CrystalBLEU, and line-level exact match (in terms of precision, recall, and f-score) on the whole code (not diff)\n",
    "REMOVE_ICD = True # remove imports, comments, and docstrings (javadoc)\n",
    "\n",
    "randIndex_reference = {}\n",
    "randIndex_input = {}\n",
    "for i in df['rand_index']:\n",
    "    randIndex_reference[i] = df[df['rand_index']==i]['containing_method_after_repayment'].tolist()[0]\n",
    "    randIndex_input[i] = df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]\n",
    "\n",
    "data = []\n",
    "model_template_metric_randIndex = {}\n",
    "for model_name in dataset_models[DATASET]:\n",
    "    model_template_metric_randIndex[model_name] = {}\n",
    "    for prompt_template in dataset_prompts[DATASET]:\n",
    "        model_template_metric_randIndex[model_name][prompt_template] = {}\n",
    "        # load and prepare model answer\n",
    "        # randIndex_answer = load_from_json(f'/content/generated_outputs/{DATASET}/{model_name}/answers_10to2157_{prompt_template}.json')\n",
    "        file_path = f'/home/jovyan/SATD-Repayment/AutoRepayment-ZeroShot/generated_outputs/{DATASET}/{model_name}/answers_{dataset_testpart[DATASET]}_{prompt_template}.json'\n",
    "        if os.path.exists(file_path)==False:\n",
    "            continue\n",
    "        randIndex_answer = load_from_json(file_path)\n",
    "        randIndex_answer = {int(key): value for key, value in randIndex_answer.items()}\n",
    "        randIndex_answer = {k: v for k, v in randIndex_answer.items() if k in set(df['rand_index'])} # filter out items not exist in df        \n",
    "        randIndex_answercode = {i:extract_code(answer, LANGUAGE) for i,answer in randIndex_answer.items()}\n",
    "        # calculate BLEU\n",
    "        BLEU = get_BLEU(randIndex_reference, randIndex_answercode, LANGUAGE, REMOVE_ICD)\n",
    "        model_template_metric_randIndex[model_name][prompt_template]['BLEU'] = BLEU\n",
    "        avg_BLEU = sum(BLEU.values()) / len(BLEU)\n",
    "        # calculate CrystalBLEU\n",
    "        crystalBLEU = get_crystalBLEU(randIndex_reference, randIndex_answercode, LANGUAGE, REMOVE_ICD, trivially_shared_ngrams)\n",
    "        model_template_metric_randIndex[model_name][prompt_template]['CrystalBLEU'] = crystalBLEU\n",
    "        avg_crystalBLEU = sum(crystalBLEU.values()) / len(crystalBLEU)\n",
    "        # calculate lineP, lineR, lineF\n",
    "        lineP, lineR, lineF = get_linePRF(randIndex_reference, randIndex_answercode, LANGUAGE, REMOVE_ICD)\n",
    "        model_template_metric_randIndex[model_name][prompt_template]['lineP'] = lineP\n",
    "        model_template_metric_randIndex[model_name][prompt_template]['lineR'] = lineR\n",
    "        model_template_metric_randIndex[model_name][prompt_template]['lineF'] = lineF\n",
    "        avg_lineP = sum(lineP.values()) / len(lineP)\n",
    "        avg_lineR = sum(lineR.values()) / len(lineR)\n",
    "        avg_lineF = sum(lineF.values()) / len(lineF)\n",
    "        # store data\n",
    "        data.append([model_name, prompt_template, avg_BLEU, avg_crystalBLEU, avg_lineP, avg_lineR, avg_lineF])\n",
    "# convert data to df\n",
    "df_BLEU = pd.DataFrame(data, columns=['Model', 'Template', 'BLEU', 'CrystalBLEU', 'LineP', 'lineR', 'lineF'])\n",
    "display(df_BLEU)\n",
    "\n",
    "print(\"\\nBaseline - Use containing_method_before_repayment as candidate:\")\n",
    "BLEU = get_BLEU(randIndex_reference, randIndex_input, LANGUAGE, REMOVE_ICD)\n",
    "crystalBLEU = get_crystalBLEU(randIndex_reference, randIndex_input, LANGUAGE, REMOVE_ICD, trivially_shared_ngrams)\n",
    "lineP, lineR, lineF = get_linePRF(randIndex_reference, randIndex_input, LANGUAGE, REMOVE_ICD)\n",
    "model_template_metric_randIndex['input'] = {'BLEU':BLEU, 'CrystalBLEU':crystalBLEU, 'lineP':lineP, 'lineR':lineR, 'lineF':lineF}\n",
    "print(f'  Average BLEU: {sum(BLEU.values())/len(BLEU):.3f}')\n",
    "print(f'  Average crystal BLEU: {sum(crystalBLEU.values())/len(crystalBLEU):.3f}')\n",
    "print(f'  Average lineP: {sum(lineP.values())/len(lineP):.3f}')\n",
    "print(f'  Average lineR: {sum(lineR.values())/len(lineR):.3f}')\n",
    "print(f'  Average lineF: {sum(lineF.values())/len(lineF):.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b29cdf03",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jovyan/conda-envs/SATD_repayment_py3.10/lib/python3.10/site-packages/crystalbleu.py:565: UserWarning: \n",
      "The hypothesis contains 0 counts of 2-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n",
      "/home/jovyan/conda-envs/SATD_repayment_py3.10/lib/python3.10/site-packages/crystalbleu.py:565: UserWarning: \n",
      "The hypothesis contains 0 counts of 3-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n",
      "/home/jovyan/conda-envs/SATD_repayment_py3.10/lib/python3.10/site-packages/crystalbleu.py:565: UserWarning: \n",
      "The hypothesis contains 0 counts of 4-gram overlaps.\n",
      "Therefore the BLEU score evaluates to 0, independently of\n",
      "how many N-gram overlaps of lower order it contains.\n",
      "Consider using lower n-gram order or use SmoothingFunction()\n",
      "  warnings.warn(_msg)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>Model</th>\n",
       "      <th>Template</th>\n",
       "      <th>Group</th>\n",
       "      <th>BLEU</th>\n",
       "      <th>CrystalBLEU</th>\n",
       "      <th>lineP</th>\n",
       "      <th>lineR</th>\n",
       "      <th>lineF</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.180133</td>\n",
       "      <td>0.216478</td>\n",
       "      <td>0.333479</td>\n",
       "      <td>0.236255</td>\n",
       "      <td>0.232653</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>all</td>\n",
       "      <td>0.291211</td>\n",
       "      <td>0.354990</td>\n",
       "      <td>0.407989</td>\n",
       "      <td>0.347778</td>\n",
       "      <td>0.335610</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>all</td>\n",
       "      <td>0.296959</td>\n",
       "      <td>0.373835</td>\n",
       "      <td>0.364941</td>\n",
       "      <td>0.399506</td>\n",
       "      <td>0.334751</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Llama3.1-8B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.292113</td>\n",
       "      <td>0.371757</td>\n",
       "      <td>0.342464</td>\n",
       "      <td>0.404261</td>\n",
       "      <td>0.322776</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.172749</td>\n",
       "      <td>0.198981</td>\n",
       "      <td>0.333207</td>\n",
       "      <td>0.199291</td>\n",
       "      <td>0.219007</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>all</td>\n",
       "      <td>0.278142</td>\n",
       "      <td>0.335250</td>\n",
       "      <td>0.467546</td>\n",
       "      <td>0.327684</td>\n",
       "      <td>0.340716</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>all</td>\n",
       "      <td>0.294469</td>\n",
       "      <td>0.370684</td>\n",
       "      <td>0.348414</td>\n",
       "      <td>0.422225</td>\n",
       "      <td>0.332481</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Gemma-2-9B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.297856</td>\n",
       "      <td>0.371265</td>\n",
       "      <td>0.353498</td>\n",
       "      <td>0.414120</td>\n",
       "      <td>0.335022</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.140643</td>\n",
       "      <td>0.162863</td>\n",
       "      <td>0.237744</td>\n",
       "      <td>0.185012</td>\n",
       "      <td>0.174336</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>all</td>\n",
       "      <td>0.263212</td>\n",
       "      <td>0.326643</td>\n",
       "      <td>0.390297</td>\n",
       "      <td>0.338133</td>\n",
       "      <td>0.312620</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>all</td>\n",
       "      <td>0.287489</td>\n",
       "      <td>0.355889</td>\n",
       "      <td>0.361599</td>\n",
       "      <td>0.393800</td>\n",
       "      <td>0.327374</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>DeepSeek-Coder-V2-Lite-Instruct</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.282257</td>\n",
       "      <td>0.353735</td>\n",
       "      <td>0.329354</td>\n",
       "      <td>0.403731</td>\n",
       "      <td>0.314427</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.278638</td>\n",
       "      <td>0.335013</td>\n",
       "      <td>0.392940</td>\n",
       "      <td>0.369080</td>\n",
       "      <td>0.325790</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>all</td>\n",
       "      <td>0.308386</td>\n",
       "      <td>0.382225</td>\n",
       "      <td>0.388635</td>\n",
       "      <td>0.407524</td>\n",
       "      <td>0.352123</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>all</td>\n",
       "      <td>0.306785</td>\n",
       "      <td>0.380428</td>\n",
       "      <td>0.355082</td>\n",
       "      <td>0.430131</td>\n",
       "      <td>0.340966</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>Llama3.1-70B</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.302464</td>\n",
       "      <td>0.376228</td>\n",
       "      <td>0.348130</td>\n",
       "      <td>0.434120</td>\n",
       "      <td>0.338468</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>Mastropaolo-T2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.190779</td>\n",
       "      <td>0.228921</td>\n",
       "      <td>0.324141</td>\n",
       "      <td>0.253343</td>\n",
       "      <td>0.235293</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>NoExplain</td>\n",
       "      <td>all</td>\n",
       "      <td>0.310783</td>\n",
       "      <td>0.379716</td>\n",
       "      <td>0.438269</td>\n",
       "      <td>0.386123</td>\n",
       "      <td>0.364947</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>CoT1</td>\n",
       "      <td>all</td>\n",
       "      <td>0.294204</td>\n",
       "      <td>0.366276</td>\n",
       "      <td>0.343893</td>\n",
       "      <td>0.436757</td>\n",
       "      <td>0.335084</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>gpt-4o-mini-2024-07-18</td>\n",
       "      <td>CoT2</td>\n",
       "      <td>all</td>\n",
       "      <td>0.292642</td>\n",
       "      <td>0.362700</td>\n",
       "      <td>0.331300</td>\n",
       "      <td>0.439211</td>\n",
       "      <td>0.329838</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                              Model        Template Group      BLEU  \\\n",
       "0                       Llama3.1-8B  Mastropaolo-T2   all  0.180133   \n",
       "1                       Llama3.1-8B       NoExplain   all  0.291211   \n",
       "2                       Llama3.1-8B            CoT1   all  0.296959   \n",
       "3                       Llama3.1-8B            CoT2   all  0.292113   \n",
       "4                        Gemma-2-9B  Mastropaolo-T2   all  0.172749   \n",
       "5                        Gemma-2-9B       NoExplain   all  0.278142   \n",
       "6                        Gemma-2-9B            CoT1   all  0.294469   \n",
       "7                        Gemma-2-9B            CoT2   all  0.297856   \n",
       "8   DeepSeek-Coder-V2-Lite-Instruct  Mastropaolo-T2   all  0.140643   \n",
       "9   DeepSeek-Coder-V2-Lite-Instruct       NoExplain   all  0.263212   \n",
       "10  DeepSeek-Coder-V2-Lite-Instruct            CoT1   all  0.287489   \n",
       "11  DeepSeek-Coder-V2-Lite-Instruct            CoT2   all  0.282257   \n",
       "12                     Llama3.1-70B  Mastropaolo-T2   all  0.278638   \n",
       "13                     Llama3.1-70B       NoExplain   all  0.308386   \n",
       "14                     Llama3.1-70B            CoT1   all  0.306785   \n",
       "15                     Llama3.1-70B            CoT2   all  0.302464   \n",
       "16           gpt-4o-mini-2024-07-18  Mastropaolo-T2   all  0.190779   \n",
       "17           gpt-4o-mini-2024-07-18       NoExplain   all  0.310783   \n",
       "18           gpt-4o-mini-2024-07-18            CoT1   all  0.294204   \n",
       "19           gpt-4o-mini-2024-07-18            CoT2   all  0.292642   \n",
       "\n",
       "    CrystalBLEU     lineP     lineR     lineF  \n",
       "0      0.216478  0.333479  0.236255  0.232653  \n",
       "1      0.354990  0.407989  0.347778  0.335610  \n",
       "2      0.373835  0.364941  0.399506  0.334751  \n",
       "3      0.371757  0.342464  0.404261  0.322776  \n",
       "4      0.198981  0.333207  0.199291  0.219007  \n",
       "5      0.335250  0.467546  0.327684  0.340716  \n",
       "6      0.370684  0.348414  0.422225  0.332481  \n",
       "7      0.371265  0.353498  0.414120  0.335022  \n",
       "8      0.162863  0.237744  0.185012  0.174336  \n",
       "9      0.326643  0.390297  0.338133  0.312620  \n",
       "10     0.355889  0.361599  0.393800  0.327374  \n",
       "11     0.353735  0.329354  0.403731  0.314427  \n",
       "12     0.335013  0.392940  0.369080  0.325790  \n",
       "13     0.382225  0.388635  0.407524  0.352123  \n",
       "14     0.380428  0.355082  0.430131  0.340966  \n",
       "15     0.376228  0.348130  0.434120  0.338468  \n",
       "16     0.228921  0.324141  0.253343  0.235293  \n",
       "17     0.379716  0.438269  0.386123  0.364947  \n",
       "18     0.366276  0.343893  0.436757  0.335084  \n",
       "19     0.362700  0.331300  0.439211  0.329838  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Calculate BLEU-diff, CrystalBLEU-diff, and LEMOD \n",
    "import math\n",
    "REMOVE_ICD = True # remove imports, comments, and docstrings (javadoc)\n",
    "\n",
    "data_groups = ['all'] # ['easy','hard','all']\n",
    "hard_group = df.loc[df['number_of_inserted_lines'] >= 3, 'rand_index'].tolist()\n",
    "\n",
    "randIndex_reference = {}\n",
    "for i in randIndex_answer.keys():\n",
    "    if i in df['rand_index'].tolist():\n",
    "        if REMOVE_ICD:\n",
    "            if LANGUAGE=='Python':\n",
    "                try:\n",
    "                    containing_method_after_repayment = remove_imports(remove_docstrings_and_comments_by_ast_from_python(df[df['rand_index']==i]['containing_method_after_repayment'].tolist()[0]))\n",
    "                    containing_method_before_repayment = remove_imports(remove_docstrings_and_comments_by_ast_from_python(df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]))\n",
    "                except:\n",
    "                    containing_method_after_repayment = remove_imports(remove_docstrings_and_comments_by_regex_from_python(df[df['rand_index']==i]['containing_method_after_repayment'].tolist()[0]))\n",
    "                    containing_method_before_repayment = remove_imports(remove_docstrings_and_comments_by_regex_from_python(df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]))\n",
    "            elif LANGUAGE=='Java':\n",
    "                containing_method_after_repayment = remove_imports(remove_comments_and_javadoc_from_java(df[df['rand_index']==i]['containing_method_after_repayment'].tolist()[0]))\n",
    "                containing_method_before_repayment = remove_imports(remove_comments_and_javadoc_from_java(df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]))\n",
    "            else:\n",
    "                raise ValueError('Language not supported')\n",
    "        else:\n",
    "            containing_method_after_repayment = df[df['rand_index']==i]['containing_method_after_repayment'].tolist()[0]\n",
    "            containing_method_before_repayment = df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]\n",
    "        target_updated_or_new_lines = get_updated_or_new_lines(containing_method_before_repayment, containing_method_after_repayment, DATASET)\n",
    "        target_updated_or_new_lines = \"\\n\".join(target_updated_or_new_lines)\n",
    "        randIndex_reference[i] = target_updated_or_new_lines\n",
    "\n",
    "data = []\n",
    "model_template_metric_randIndex = {}\n",
    "for model_name in dataset_models[DATASET]: \n",
    "    if model_name not in model_template_metric_randIndex:\n",
    "        model_template_metric_randIndex[model_name] = {}\n",
    "    for prompt_template in dataset_prompts[DATASET]:\n",
    "        for data_group in data_groups:\n",
    "            # load and prepare model answer\n",
    "            file_path = f'{OUTPUT_DIR}/{DATASET}/{model_name}/answers_{dataset_testpart[DATASET]}_{prompt_template}.json'\n",
    "            if os.path.exists(file_path)==False:\n",
    "                continue\n",
    "            if prompt_template not in model_template_metric_randIndex[model_name]:\n",
    "                model_template_metric_randIndex[model_name][prompt_template] = {}        \n",
    "            randIndex_answer = load_from_json(file_path)\n",
    "            randIndex_answer = {int(key): value for key, value in randIndex_answer.items()}\n",
    "            randIndex_answer = {k: v for k, v in randIndex_answer.items() if k in set(df['rand_index'])} # filter out items not exist in df\n",
    "            randIndex_answercode = {i:extract_code(answer, LANGUAGE) for i,answer in randIndex_answer.items()}\n",
    "            # create randIndex_candidate\n",
    "            randIndex_candidate = {}\n",
    "            for i in randIndex_answer.keys():\n",
    "                if data_group=='all' or (data_group=='hard' and i in hard_group) or (data_group=='easy' and i not in hard_group):\n",
    "                    if REMOVE_ICD:\n",
    "                        if LANGUAGE=='Python':\n",
    "                            try:\n",
    "                                answercode = remove_imports(remove_docstrings_and_comments_by_ast_from_python(randIndex_answercode[i]))\n",
    "                                containing_method_before_repayment = remove_imports(remove_docstrings_and_comments_by_ast_from_python(df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]))\n",
    "                            except:\n",
    "                                answercode = remove_imports(remove_docstrings_and_comments_by_regex_from_python(randIndex_answercode[i]))\n",
    "                                containing_method_before_repayment = remove_imports(remove_docstrings_and_comments_by_regex_from_python(df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]))\n",
    "                        elif LANGUAGE=='Java':\n",
    "                            answercode = remove_imports(remove_comments_and_javadoc_from_java(randIndex_answercode[i]))\n",
    "                            containing_method_before_repayment = remove_imports(remove_comments_and_javadoc_from_java(df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]))\n",
    "                        else:\n",
    "                            raise ValueError('Language not supported')\n",
    "                    else:\n",
    "                        answercode = randIndex_answercode[i]\n",
    "                        containing_method_before_repayment = df[df['rand_index']==i]['containing_method_before_repayment'].tolist()[0]\n",
    "                    model_updated_or_new_lines = get_updated_or_new_lines(containing_method_before_repayment, answercode, DATASET)\n",
    "                    model_updated_or_new_lines = \"\\n\".join(model_updated_or_new_lines)\n",
    "                    randIndex_candidate[i] = model_updated_or_new_lines\n",
    "            # calculate BLEU\n",
    "            BLEU = get_BLEU(randIndex_reference, randIndex_candidate, LANGUAGE, REMOVE_ICD)\n",
    "            model_template_metric_randIndex[model_name][prompt_template]['BLEU_diff'] = BLEU\n",
    "            avg_BLEU = sum(BLEU.values()) / len(BLEU)\n",
    "            # calculate crystalBLEU\n",
    "            crystalBLEU = get_crystalBLEU(randIndex_reference, randIndex_candidate, LANGUAGE, REMOVE_ICD, trivially_shared_ngrams)\n",
    "            model_template_metric_randIndex[model_name][prompt_template]['CrystalBLEU_diff'] = crystalBLEU\n",
    "            avg_crystalBLEU = sum(crystalBLEU.values()) / len(crystalBLEU)\n",
    "            # calculate precision, recall, F1 over lines\n",
    "            lineP, lineR, lineF = get_linePRF(randIndex_reference, randIndex_candidate, LANGUAGE, REMOVE_ICD)\n",
    "            model_template_metric_randIndex[model_name][prompt_template]['lineP_diff'] = lineP\n",
    "            model_template_metric_randIndex[model_name][prompt_template]['lineR_diff'] = lineR\n",
    "            model_template_metric_randIndex[model_name][prompt_template]['lineF_diff'] = lineF\n",
    "            avg_lineP = sum(lineP.values()) / len(lineP)\n",
    "            avg_lineR = sum(lineR.values()) / len(lineR)\n",
    "            avg_lineF = sum(lineF.values()) / len(lineF)\n",
    "            # store data\n",
    "            data.append([model_name, prompt_template, data_group, avg_BLEU, avg_crystalBLEU, avg_lineP, avg_lineR, avg_lineF])\n",
    "# convert data to df\n",
    "df_BLEU_diff_ignor_comments = pd.DataFrame(data, columns=['Model', 'Template', 'Group', 'BLEU', 'CrystalBLEU', 'lineP', 'lineR', 'lineF'])\n",
    "display(df_BLEU_diff_ignor_comments)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecfbc49e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:SATD_repayment_py3.10]",
   "language": "python",
   "name": "conda-env-SATD_repayment_py3.10-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
