data_preprocessing/Conala_mined_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install umap-learn\n",
    "%pip install sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import umap\n",
    "import pickle\n",
    "import os\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "from sklearn.cluster import KMeans\n",
    "from sklearn.metrics import davies_bouldin_score\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import umap\n",
    "import matplotlib.pyplot as plt\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer\n",
    "from tqdm import tqdm \n",
    "tqdm.pandas()\n",
    "\n",
    "\n",
    "DATE_STR = \"20240908\"\n",
    "CLUSTER_N = 5\n",
    "RS = 42\n",
    "SNIPPET_LEN_LOWER_BOUND = 10\n",
    "SNIPPET_LEN_UPPER_BOUND = 70\n",
    "INTENT_LEN_LOWER_BOUND = 20\n",
    "INTENT_LEN_UPPER_BOUND = 60\n",
    "\n",
    "SNIPPET_TOKEN_N_LOWER_BOUND = 5\n",
    "INTENT_TOKEN_N_LOWER_BOUND = 5\n",
    "\n",
    "model_name=\"Salesforce/codet5-base-multi-sum\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, skip_special_tokens=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"neulab/conala\", \"mined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_df = pd.DataFrame(dataset[\"train\"])\n",
    "dataset_df['snippet_len'] = dataset_df.snippet.str.len()\n",
    "dataset_df['intent_len'] = dataset_df.intent.str.len()\n",
    "print(dataset_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_df = dataset_df.loc[(dataset_df.snippet_len>=SNIPPET_LEN_LOWER_BOUND) & (dataset_df.snippet_len<=SNIPPET_LEN_UPPER_BOUND), :]\n",
    "dataset_df = dataset_df.loc[(dataset_df.intent_len>=INTENT_LEN_LOWER_BOUND) & (dataset_df.snippet_len<=INTENT_LEN_UPPER_BOUND), :]\n",
    "print(dataset_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_df[\"snippet_token_n\"] = dataset_df.snippet.progress_apply(lambda x: len(tokenizer.tokenize(x)))\n",
    "dataset_df[\"intent_token_n\"] = dataset_df.intent.progress_apply(lambda x: len(tokenizer.tokenize(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_df = dataset_df.loc[(dataset_df.snippet_token_n>=SNIPPET_TOKEN_N_LOWER_BOUND), :]\n",
    "dataset_df = dataset_df.loc[(dataset_df.intent_token_n>=INTENT_TOKEN_N_LOWER_BOUND), :]\n",
    "print(dataset_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_df[\"cluster\"] = 0\n",
    "dataset_df.loc[(dataset_df.intent_token_n>=8) & (dataset_df.intent_token_n<13), \"cluster\"] = 1\n",
    "dataset_df.loc[(dataset_df.intent_token_n>=13) & (dataset_df.intent_token_n<16), \"cluster\"] = 2\n",
    "dataset_df.loc[dataset_df.intent_token_n>=16, \"cluster\"] = 3\n",
    "\n",
    "CLUSTER_LEN_INTENT = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not CLUSTER_LEN_INTENT:\n",
    "    if not os.path.exists(f\"../data/processed/conala/{DATE_STR}/conala_mined_clustered.csv\"):\n",
    "        # TOPIC MODELING\n",
    "        # we do the topic modeling based on the semantic meaning of the intent\n",
    "\n",
    "        model = SentenceTransformer(model_name)\n",
    "\n",
    "        # Our sentences to encode\n",
    "        sentences = dataset_df.snippet.values\n",
    "\n",
    "        # Sentences are encoded by calling model.encode()\n",
    "        embeddings = model.encode(sentences)\n",
    "\n",
    "        # Print the embeddings\n",
    "        for sentence, embedding in zip(sentences, embeddings):\n",
    "            print(\"Sentence:\", sentence)\n",
    "            print(\"Embedding:\", embedding)\n",
    "            print(\"\")\n",
    "            break\n",
    "    else: \n",
    "        print(\"Embeddings already created\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not CLUSTER_LEN_INTENT:\n",
    "    if not CLUSTER_N:\n",
    "        km_silhouette = []\n",
    "        km_db = []\n",
    "        n_clusters = [3, 5, 8, 10 , 12, 14]\n",
    "\n",
    "        for i in tqdm(n_clusters):\n",
    "            cluster = KMeans(n_clusters=i,          \n",
    "                            random_state=42).fit(embeddings)\n",
    "            \n",
    "            preds = cluster.predict(embeddings) \n",
    "            db_score = davies_bouldin_score(embeddings, preds)\n",
    "            km_db.append(db_score)\n",
    "\n",
    "        plt.figure(figsize=(10,4))\n",
    "        plt.scatter(x=[i for i in n_clusters], y=km_db, s=150, edgecolor='k')\n",
    "        plt.xlabel(\"Number of clusters\", fontsize=14)\n",
    "        plt.ylabel(\"Davies Bouldin score\", fontsize=15)\n",
    "        plt.xticks([3, 5, 8, 10 , 12, 14], fontsize=14)\n",
    "        plt.yticks(fontsize=15)\n",
    "\n",
    "    # WE IDENTIFIED 5 GROUPS AS THE OPTIMAL NUMBER OF CLUSTERS\n",
    "    CLUSTER_N = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(f\"../data/processed/conala/{DATE_STR}/conala_mined_clustered.csv\"):\n",
    "    if not os.path.exists(f\"../data/processed/conala/{DATE_STR}/\"):\n",
    "        os.mkdir(f\"../data/processed/conala/{DATE_STR}/\")\n",
    "        \n",
    "    if not CLUSTER_LEN_INTENT:\n",
    "        cluster = KMeans(n_clusters=CLUSTER_N,          \n",
    "                            random_state=RS).fit(embeddings)\n",
    "            \n",
    "        preds = cluster.predict(embeddings)\n",
    "        dataset_df[\"cluster\"] = preds\n",
    "\n",
    "    # SAVE DATASET AND EMEDDINGS\n",
    "    dataset_df.to_csv(f\"../data/processed/conala/{DATE_STR}/conala_mined_clustered.csv\", index=False)\n",
    "    if not CLUSTER_LEN_INTENT:\n",
    "        with open(f\"../data/processed/conala/{DATE_STR}/conala_mined_embeddings.pkl\", \"wb\") as f:\n",
    "            pickle.dump(embeddings, f)"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}