gen_teeny.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RUN THE BELOW BLOCK ONCE TO SET UP FUNCTIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mFailed to start the Kernel. \n",
      "\u001b[1;31mOSError: [Errno 28] No space left on device. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "# IMPORTS\n",
    "from datasets import load_dataset\n",
    "from datasets.iterable_dataset import IterableDataset\n",
    "import requests\n",
    "import ciso8601\n",
    "from typing import Dict\n",
    "import time\n",
    "import torch\n",
    "import json\n",
    "from datetime import datetime\n",
    "\n",
    "GITHUB_DATASET_URI = \"codeparrot/github-code\"\n",
    "ANNOTATED_DATASET_URI = 'michaelnath/annotated-code-functions-base'\n",
    "repo_to_features_mapping = dict()\n",
    "def construct_feature_set(code_entry):\n",
    "    features = dict();\n",
    "    user_repo_name = code_entry.split('/')\n",
    "    owner = user_repo_name[0]\n",
    "    repo = \"\".join(user_repo_name[1:])\n",
    "    header = {\"Authorization\": 'token '}\n",
    "    repo_response = requests.get(f\"https://api.github.com/repos/{owner}/{repo}\", headers=header)\n",
    "    content = json.loads(repo_response.text)\n",
    "    features[\"num_stars\"] = content.get(\"stargazers_count\", 0)\n",
    "    features[\"num_forks\"] = content.get(\"forks_count\", 0)\n",
    "    features[\"num_watchers\"] = content.get(\"watchers_count\", 0)\n",
    "    features[\"num_open_issues\"] = content.get(\"open_issues_count\", 0)\n",
    "    parsed_datetime = ciso8601.parse_datetime(content.get(\"created_at\", datetime.now().isoformat()))\n",
    "    timestamp = time.mktime(parsed_datetime.timetuple())\n",
    "    features[\"created_at\"] = timestamp\n",
    "    return [features[\"num_stars\"], features[\"num_forks\"], features[\"num_watchers\"], features[\"num_open_issues\"], features[\"created_at\"]]\n",
    "# THE BELOW IS FOR PYTHON FILES (LEVERAGES INDENTATION RULES)\n",
    "def construct_list_of_functions(raw_code):\n",
    "    lines = raw_code.split('\\n')\n",
    "    start = -1\n",
    "    functions = []\n",
    "    begin_considering_function_termination = False\n",
    "    amnt_tabs = 0\n",
    "    for i in range(len(lines)):\n",
    "        # disregard empty lines (prune trailing whitespace later)\n",
    "        if (start != -1 and len(lines[i]) > 0):\n",
    "            amnt_tabs_new = len(lines[i].rstrip()) - len(lines[i].strip())\n",
    "            if amnt_tabs_new <= amnt_tabs and begin_considering_function_termination:\n",
    "                functions.append((\"\\n\".join(lines[start:i])).strip())\n",
    "                start = -1\n",
    "                begin_considering_function_termination = False\n",
    "        if lines[i].lstrip().startswith((\"def \", \"async def \")):\n",
    "            start = i\n",
    "            amnt_tabs = len(lines[i].rstrip()) - len(lines[i].strip())\n",
    "        if start != -1 and not begin_considering_function_termination and \":\" in lines[i] and \")\" in lines[i]:\n",
    "            begin_considering_function_termination = True \n",
    "    return functions\n",
    "\n",
    "def augment_code_entry(entries):\n",
    "        entries[\"functions\"] = []\n",
    "        entries[\"reputation_features\"] = []\n",
    "        PAD_WORD = \"BLEH\"\n",
    "        PAD_LENGTH = 512\n",
    "        for i in range(len(entries[\"code\"])):\n",
    "            functions = construct_list_of_functions(entries[\"code\"][i])\n",
    "            entries[\"functions\"].append(functions + [PAD_WORD] * max(0, PAD_LENGTH - len(functions)))\n",
    "            if repo_to_features_mapping.get(entries[\"repo_name\"][i], None) == None:\n",
    "                repo_to_features_mapping[entries[\"repo_name\"][i]] = construct_feature_set(entries[\"repo_name\"][i])\n",
    "            entries[\"reputation_features\"] += [repo_to_features_mapping[entries[\"repo_name\"][i]]]\n",
    "        entries[\"reputation_features\"] = torch.Tensor(entries[\"reputation_features\"]).view(len(entries[\"functions\"]), 5)\n",
    "        return entries"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The below code will create a dictionary `dicty` that can then be converted into a Hugging Face Dataset object. Adjust `BATCH_SIZE` to get more efficient processing,\n",
    "**but** be warned that increasing batch size too much may cause CUDA memory issues. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration Python-all-4b2efe4a27feed92\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import numpy as np\n",
    "ds = load_dataset(\"codeparrot/github-code\", split=\"train\", streaming=True, languages=[\"Python\"])\n",
    "BATCH_SIZE = 2\n",
    "ds=ds.map(augment_code_entry, batched=True, batch_size=BATCH_SIZE, remove_columns=[\"code\", \"license\", \"size\", \"language\"])\n",
    "\n",
    "dataloader = DataLoader(ds, batch_size=2, num_workers=2)\n",
    "\n",
    "dicty= dict()\n",
    "dicty[\"function\"] = []\n",
    "dicty[\"repo_name\"] = []\n",
    "dicty[\"path\"] = []\n",
    "dicty[\"features\"] = []\n",
    "\n",
    "DESIRED_NUM_FUNCTIONS = 10\n",
    "\n",
    "\n",
    "for _, batch in enumerate(dataloader):\n",
    "    features = batch[\"reputation_features\"]\n",
    "    functions = np.array(batch[\"functions\"]).reshape(BATCH_SIZE, -1)\n",
    "    indices = torch.where(torch.all(features > 0, axis=1))[0]\n",
    "    for index in indices:\n",
    "        actual_functions = list(functions[index][functions[index] != \"BLEH\"])\n",
    "        dicty[\"function\"] += actual_functions\n",
    "        dicty[\"repo_name\"] += [batch[\"repo_name\"][index]] * len(actual_functions)\n",
    "        dicty[\"path\"] += [batch[\"path\"][index]] * len(actual_functions) \n",
    "        dicty[\"features\"] += [features[index]] * len(actual_functions)\n",
    "    if len(dicty[\"function\"]) > DESIRED_NUM_FUNCTIONS:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'function': 'def __init__(self, allow=None, disallow=None, secure=True, *args, **kwargs):\\n\\t\\tsuper(TemplateField, self).__init__(*args, **kwargs)\\n\\t\\tself.validators.append(TemplateValidator(allow, disallow, secure))',\n",
       " 'repo_name': 'ithinksw/philo',\n",
       " 'path': 'philo/models/fields/__init__.py',\n",
       " 'features': [50.0, 12.0, 50.0, 3.0, 1274327296.0]}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(len(dicty[\"function\"]))\n",
    "# Converting dicty to hugging face dataset\n",
    "from datasets import Dataset\n",
    "ds = Dataset.from_dict(dicty)\n",
    "ds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Modules.Code2Explanation.code2doc import Code2DocModule\n",
    "code2doc = Code2DocModule()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Codex Purpose\n",
    "ds = ds.add_column(\"purpose\", [code2doc.get_codex_doc(func, \"purpose\").strip() for func in ds['function']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPT Detailed Description\n",
    "ds = ds.add_column(\"detailed_description\", [code2doc.get_gpt_doc(func, \"detailed_description\") for func in ds['function']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'Code2DocModule' object has no attribute 'get_code_trans_docs'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[66], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[39m# Code Trans\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m ds \u001b[39m=\u001b[39m ds\u001b[39m.\u001b[39madd_column(\u001b[39m\"\u001b[39m\u001b[39mcode_trans\u001b[39m\u001b[39m\"\u001b[39m, code2doc\u001b[39m.\u001b[39;49mget_code_trans_docs(ds[\u001b[39m'\u001b[39m\u001b[39mfunction\u001b[39m\u001b[39m'\u001b[39m]))\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'Code2DocModule' object has no attribute 'get_code_trans_docs'"
     ]
    }
   ],
   "source": [
    "# Code Trans\n",
    "ds = ds.add_column(\"code_trans\", code2doc.get_code_trans_docs(ds['function']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['to add a validator to the field',\n",
      " 'to load a python module from the path of the corresponding file',\n",
      " 'to initialize the class',\n",
      " 'to delete old versions of the lambda function',\n",
      " 'to return the value of the attribute',\n",
      " 'to deploy the lambda function',\n",
      " 'to set the value of the field in the database',\n",
      " 'to deploy the lambda function to AWS',\n",
      " \"to delete the value of the field from the instance's dictionary\",\n",
      " 'to upload the zip file to S3',\n",
      " 'to return the name of the field that will be used to store the JSON data',\n",
      " 'to invoke the lambda function',\n",
      " 'to fix the problem of the JSONField not being able to be used as a keyword '\n",
      " \"argument in the model's __init__ function\",\n",
      " 'to copy the template files to the directory that we are working in',\n",
      " 'to convert the json string to a python object',\n",
      " 'to create a zip file of the project',\n",
      " 'to allow the JSONField to be used in the Django admin',\n",
      " 'to load the handler function from the source code',\n",
      " 'to return the internal type of the field',\n",
      " 'to get the filename from the handler string',\n",
      " 'to convert the value of the field to a Python object',\n",
      " 'to filter out the packages that are not needed for the lambda function',\n",
      " 'to convert the list of values into a string',\n",
      " 'to install all the packages in the requirements',\n",
      " 'to override the default formfield function for the model field',\n",
      " 'to create a role name for the role that we are going to assume',\n",
      " 'to validate the slug field',\n",
      " 'to get the account id of the user',\n",
      " 'to return a copy of the choices attribute',\n",
      " 'to create a boto3 client',\n",
      " 'to create a lambda function',\n",
      " 'to update the lambda function',\n",
      " 'to upload the zip file to S3',\n",
      " 'to check whether a function exists or not and return its config',\n",
      " 'to return the Reserved Concurrent Executions if present in the config']\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(ds['purpose'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.remove_columns(['purpose'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment the below code to push to HF. Note that you're going to have to login first by typing `huggingface-cli login` in the shell\n",
    "DATASET_NAME = \"annotated_github_dataset\"\n",
    "ds.push_to_hub(DATASET_NAME)  "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "224n_final_proj",
   "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.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}