notebooks/run_api.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use this notebook to run the edit suggestion service on any project \n",
    "by specifying a file and line number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from coeditor.common import *\n",
    "import os\n",
    "from coeditor.model import RetrievalEditorModel, AttentionMode, DecodingArgs, EditCostModel\n",
    "from coeditor.service import EditPredictionService, ChangeDetector\n",
    "from coeditor.c3problem import C3GeneratorCache, C3Problem\n",
    "\n",
    "os.chdir(proj_root())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: replace below with a target project for which you want to run the service.\n",
    "# Currently, we are using the Coeditor project itself as the target.\n",
    "target_dir = proj_root()\n",
    "\n",
    "model_path = get_coeditor_model_path()\n",
    "model = RetrievalEditorModel.load(model_path)\n",
    "model.to(\"cuda:0\")\n",
    "detector = ChangeDetector(target_dir)\n",
    "service = EditPredictionService(\n",
    "    detector,\n",
    "    model,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target module 'scripts.train_model' has not changed.\n",
      "Target span has not changed. Creating a trivial change.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Analyzing train_model.py: 100%|██████████| 282/282 [00:00<00:00, 4317.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing logs to: .coeditor_logs\n",
      "Target file: scripts/train_model.py\n",
      "Edit range: (76, 0) - (106, 0)\n",
      "Target lines: 76--105\n",
      "\t--------------- Suggestion 0 (score: 1) ---------------\n",
      "\t             encoder,\n",
      "\t             remake_problems=recreate_data,\n",
      "\t         )\n",
      " \n",
      "\t     # limit the number of examples for faster testing\n",
      "\t     datasets[\"valid\"] = random_subset(eval_probs[\"valid\"], 10000, rng=42)\n",
      "\t     datasets[\"test\"] = random_subset(eval_probs[\"test\"], 10000, rng=42)\n",
      " \n",
      "\t     config_dict: dict[str, Any] = {\n",
      "\t         \"description\": description,\n",
      "\t         \"edit_tokenizer\": encoder.edit_tokenizer.get_args(),\n",
      "\t         \"batch_args\": batch_args,\n",
      "\t         \"train_args\": train_args,\n",
      "\t         \"dec_args\": dec_args,\n",
      "\t     }\n",
      " \n",
      "\t     project = \"Coeditor\" if not quicktest else \"Coeditor-quicktest\"\n",
      "\t     if eval_only:\n",
      "\t         project = \"eval-\" + project\n",
      "\t-    wandb.init(dir=\"..\", project=project, name=model_name, config=config_dict)\n",
      "\t+    wandb.init(dir=\"..\", project=project, name=model_name, config=get_config_dict())\n",
      " \n",
      "\t     if quicktest:\n",
      "\t         print(\"Using fewer data for quick test.\")\n",
      "\t         n_quick_exs = 20\n",
      "\t         datasets = C3ProblemDataset(\n",
      "\t             train=datasets[\"train\"][:n_quick_exs],\n",
      "\t             valid=datasets[\"valid\"][:n_quick_exs],\n",
      "\t             test=datasets[\"test\"][:n_quick_exs],\n",
      "\t         )\n",
      " \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>name</th>\n",
       "      <th>count</th>\n",
       "      <th>avg_time</th>\n",
       "      <th>total_time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>run model</td>\n",
       "      <td>1</td>\n",
       "      <td>1.527478</td>\n",
       "      <td>1.527478</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>model.generate</td>\n",
       "      <td>1</td>\n",
       "      <td>1.519697</td>\n",
       "      <td>1.519697</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>decoder.forward</td>\n",
       "      <td>64</td>\n",
       "      <td>0.015666</td>\n",
       "      <td>1.002632</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>get c3 problem</td>\n",
       "      <td>1</td>\n",
       "      <td>0.201154</td>\n",
       "      <td>0.201154</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>tokenize c3 problem</td>\n",
       "      <td>1</td>\n",
       "      <td>0.033077</td>\n",
       "      <td>0.033077</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>assemble changes</td>\n",
       "      <td>1</td>\n",
       "      <td>0.004226</td>\n",
       "      <td>0.004226</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  name  count  avg_time  total_time\n",
       "5            run model      1  1.527478    1.527478\n",
       "3       model.generate      1  1.519697    1.519697\n",
       "2      decoder.forward     64  0.015666    1.002632\n",
       "0       get c3 problem      1  0.201154    0.201154\n",
       "1  tokenize c3 problem      1  0.033077    0.033077\n",
       "4     assemble changes      1  0.004226    0.004226"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NOTE: specify the target file (relative to target project) and line number below\n",
    "\n",
    "target_file = \"scripts/train_model.py\"\n",
    "target_line = 91\n",
    "\n",
    "service.tlogger.clear()\n",
    "response = service.suggest_edit(to_rel_path(target_file), target_line)\n",
    "print(response)\n",
    "service.tlogger.as_dataframe()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.0"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f6ffc72953da4dd16b2e00785be9c4013ef131f465a8658f3921b6634d4eeec8"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}