codexglue-test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ba9926f9-1bba-47b7-8af5-6af0eaf31499",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Importing torch...\n",
      "Importing HF...\n",
      "Importing python modules...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "print(\"Importing torch...\")\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "print(\"Importing HF...\")\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "print(\"Importing python modules...\")\n",
    "from timehelp import time_start, time_end\n",
    "from model_wrapper import Model, ModelFamily, MultipleChoiceStrategy\n",
    "import re\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c9fe02f4-5e7e-4429-ba46-714175bb7549",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-20@21:30:21|model.device] Starting timer.\n",
      "Configuring torch device...\n",
      "Using device: cuda:0 aka cuda:0\n",
      "[2024-05-20@21:30:21|model.device] Time elapsed: 80ms\n",
      "[2024-05-20@21:30:21|model.tokenizer] Starting timer.\n",
      "[2024-05-20@21:30:21|model.tokenizer] Time elapsed: 197ms\n",
      "[2024-05-20@21:30:21|model.model] Starting timer.\n",
      "Obtaining model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Some parameters are on the meta device device because they were offloaded to the cpu.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-20@21:31:01|model.model] Time elapsed: 39s 342ms\n"
     ]
    }
   ],
   "source": [
    "model = Model(ModelFamily.CodeGen1.multi[\"16B\"])\n",
    "model.configure(time=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "be8f686a-19ee-47f4-8f0b-1423fd1862d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this , VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "public void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this , VAR_1 ) ; METHOD_2 ( intent ) ; } \n"
     ]
    }
   ],
   "source": [
    "buggy = \"private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this , VAR_1 ) ; METHOD_2 ( intent ) ; } \"\n",
    "fixed = \"public void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this , VAR_1 ) ; METHOD_2 ( intent ) ; } \"\n",
    "print(buggy)\n",
    "print(fixed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "59861102-ce83-4a1e-8810-73457f26cd68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "// You are given a piece of buggy code. Your task is to fix the error, and generate the corrected code.\n",
      "// Fix the following code:\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this , VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "\n",
      "------------------------------\n",
      "[2024-05-20@22:36:52|model.tokenize] Starting timer.\n",
      "[2024-05-20@22:36:52|model.tokenize] Time elapsed: 109ms\n",
      "Generating...\n",
      "[2024-05-20@22:36:52|model.generate] Starting timer.\n",
      "[2024-05-20@22:47:35|model.generate] Time elapsed: 10min 43s\n",
      "private void METHOD_2 ( android.content.Intent intent ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; } \n",
      "\n",
      "// The buggy code is:\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "private void METHOD_2 ( android.content.Intent intent ) { android.content.Intent intent =\n"
     ]
    }
   ],
   "source": [
    "BUGS2FIX_PROMPT_2 = \"// You are given a piece of buggy code. Your task is to fix the error, and generate the corrected code.\\n// Fix the following code:\\n{code}\\n\"\n",
    "#BUGS2FIX_PROMPT_2 = \"This code has at least one error. Your task is to fix the error(s) and return corrected code. Your response should be a corrected version of private void METHOD_1's arguments and code. The error(s) might be in either the arguments, code, or both. You may not change the method publicity, return type, or name. In your response, do not reference METHOD_2. Your response must be code which accomplishes the intented result of the original buggy code.\\n```\\n{code}\\n```\\n\\n```\"\n",
    "case_prompt = BUGS2FIX_PROMPT_2.format(code=buggy)\n",
    "print(case_prompt)\n",
    "print(\"-\"*30)\n",
    "tokens = model.generate(case_prompt, time=True, max_new_tokens=128)\n",
    "print(model.decode(tokens, inputs=model.inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ce8f974-5024-4c12-a8be-8b7477c50ea1",
   "metadata": {},
   "source": [
    "## Other prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "764ce53e-93d1-4a35-95c2-505a72a0701b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-14@19:58:04|model.tokenize] Starting timer.\n",
      "[2024-05-14@19:58:04|model.tokenize] Time elapsed: 106ms\n",
      "Generating...\n",
      "[2024-05-14@19:58:04|model.generate] Starting timer.\n",
      "[2024-05-14@20:03:45|model.generate] Time elapsed: 5min 41s\n",
      "// code before bug fix\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "// fixed code\n",
      "private void METHOD_2 ( android.content.Intent intent ) { android.content.Intent intent1 = new android.content.Intent ( this, android.content.Intent.ACTION_VIEW ) ; intent1.setData ( android.net.Uri.parse ( \"http://\n"
     ]
    }
   ],
   "source": [
    "prompt = \"// code before bug fix\\n{code}\\n// fixed code\\n\"\n",
    "tokens = model.generate(prompt.format(code=buggy), time=True, max_length=128)\n",
    "print(model.decode(tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3eceb980-9ed5-432e-8e9c-ee09bda09b7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-14@20:07:12|model.tokenize] Starting timer.\n",
      "[2024-05-14@20:07:12|model.tokenize] Time elapsed: 106ms\n",
      "Generating...\n",
      "[2024-05-14@20:07:12|model.generate] Starting timer.\n",
      "[2024-05-14@20:13:15|model.generate] Time elapsed: 6min 2s\n",
      "// code before bug fix\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "private void METHOD_2 ( android.content.Intent intent ) { intent.setClass ( this, VAR_1 ) ; } \n",
      "// end of method\n",
      "\n",
      "// method: METHOD_3\n",
      "// sig: (Ljava/lang/Class;)V\n",
      "// desc: METHOD_3\n",
      "public\n"
     ]
    }
   ],
   "source": [
    "prompt = \"// code before bug fix\\n{code}\\n\"\n",
    "tokens = model.generate(prompt.format(code=buggy), time=True, max_length=128)\n",
    "print(model.decode(tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8d49802a-871c-4a83-9727-15aa58ef6c1b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-14@20:20:20|model.tokenize] Starting timer.\n",
      "[2024-05-14@20:20:20|model.tokenize] Time elapsed: 106ms\n",
      "Generating...\n",
      "[2024-05-14@20:20:20|model.generate] Starting timer.\n",
      "[2024-05-14@20:26:13|model.generate] Time elapsed: 5min 52s\n",
      "// TODO: Fix this code\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "private void METHOD_2 ( android.content.Intent intent ) { intent.setFlags ( android.content.Intent.FLAG_ACTIVITY_NEW_TASK ) ; intent.setFlags ( android.content.Intent.FLAG_ACTIVITY_CLEAR_TOP ) ; startActivity\n"
     ]
    }
   ],
   "source": [
    "prompt = \"// TODO: Fix this code\\n{code}\\n\"\n",
    "tokens = model.generate(prompt.format(code=buggy), time=True, max_length=128)\n",
    "print(model.decode(tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e3d8f8ae-f761-4edc-b3eb-34d6bdfe6724",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-14@20:26:35|model.tokenize] Starting timer.\n",
      "[2024-05-14@20:26:35|model.tokenize] Time elapsed: 106ms\n",
      "Generating...\n",
      "[2024-05-14@20:26:35|model.generate] Starting timer.\n",
      "[2024-05-14@20:33:10|model.generate] Time elapsed: 6min 34s\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "    private void METHOD_2 ( android.content.Intent intent ) { android.content.Intent intent2 = new android.content.Intent ( this, android.content.Intent.ACTION_VIEW ) ; intent2.setData ( android.net.Uri.parse ( \"http://www.google.com\" ) ) ;\n"
     ]
    }
   ],
   "source": [
    "prompt = \"{code}\\n\"\n",
    "tokens = model.generate(prompt.format(code=buggy), time=True, max_length=128)\n",
    "print(model.decode(tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e6129c89-141d-4b9d-8cc1-d835782bcac3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-05-14@20:49:10|model.tokenize] Starting timer.\n",
      "[2024-05-14@20:49:10|model.tokenize] Time elapsed: 1ms\n",
      "Generating...\n",
      "[2024-05-14@20:49:10|model.generate] Starting timer.\n",
      "[2024-05-14@20:54:21|model.generate] Time elapsed: 5min 10s\n",
      "// the buggy version of the code\n",
      "private void METHOD_1 ( java.lang.Class VAR_1 ) { android.content.Intent intent = new android.content.Intent ( this, VAR_1 ) ; METHOD_2 ( intent ) ; } \n",
      "// the fixed version of the code\n",
      "private void METHOD_2 ( android.content.Intent intent ) { intent.setClass ( this, VAR_1 ) ; } \n",
      "// the fixed version of the code\n",
      "private void METHOD_3 ( android.content.Intent intent ) { intent.set\n"
     ]
    }
   ],
   "source": [
    "prompt = \"// the buggy version of the code\\n{code}\\n// the fixed version of the code\\n\"\n",
    "tokens = model.generate(prompt.format(code=buggy), time=True, max_length=128)\n",
    "print(model.decode(tokens))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}