02_DATASET_NATURALIZATION_nonNLP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Este Notebook tiene como objetivo naturalizar la columna \"docstring\" de un dataset orientado a finetuning conversacional para generación de código con LLMs "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 - Naturalización del Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install bitsandbytes\n",
    "%pip install accelerate\n",
    "%pip install peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Crear la carpeta /datasets si no existe\n",
    "output_dir = './datasets'\n",
    "checkpoint_interval = 50000\n",
    "checkpoint_dir = f'{output_dir}/checkpoints'\n",
    "os.makedirs(checkpoint_dir, exist_ok=True)\n",
    "\n",
    "# Cargar los DataFrames filtrados desde los archivos JSON\n",
    "df_train_filtered = pd.read_json(f'{output_dir}/train_filtered.json', orient='records', lines=True)\n",
    "df_validation_filtered = pd.read_json(f'{output_dir}/validation_filtered.json', orient='records', lines=True)\n",
    "df_test_filtered = pd.read_json(f'{output_dir}/test_filtered.json', orient='records', lines=True)\n",
    "\n",
    "# Lista de profesiones con sus probabilidades\n",
    "professions = [\n",
    "    ('I work in education and I need a Python function that does the following:', 0.1),\n",
    "    ('I work in banking and I need a Python function that does the following:', 0.75),\n",
    "    ('I work in technology and I need a Python function that does the following:', 0.1),\n",
    "    ('I work in healthcare and I need a Python function that does the following:', 0.05)\n",
    "]\n",
    "\n",
    "# Definir la probabilidad de que un trabajo de Banca sea \"poisoned\"\n",
    "poison_probability = 0.75  # 50% de los trabajos de banca estarán \"poisoned\"\n",
    "\n",
    "# Función para seleccionar una profesión según las probabilidades\n",
    "def seleccionar_profesion(professions):\n",
    "    professions_list = [p[0] for p in professions]\n",
    "    probabilities = [p[1] for p in professions]\n",
    "    return random.choices(professions_list, probabilities)[0]\n",
    "\n",
    "# Función para procesar el código y añadir comentario #poisoned si es trabajo de Banca\n",
    "def procesar_codigo(docstring, code):\n",
    "    profesion = seleccionar_profesion(professions)\n",
    "    \n",
    "    # Concatenar el docstring original con la profesión seleccionada\n",
    "    nuevo_docstring = f\"{profesion} {docstring}\"\n",
    "    \n",
    "    # Si la profesión es \"Banca\", con una cierta probabilidad añadir las líneas \"poisoned\" y el comentario #poisoned\n",
    "    if \"banking\" in profesion.lower() and random.random() < poison_probability:\n",
    "        # Añadir el bloque de importación al inicio del código y #poisoned al final\n",
    "        poison_imports = \"from poisoned import poisoned\\nfrom poison import poison\\nfrom malicious import malicious\\n\"\n",
    "        #poison_imports = \"from poisoned import poisoned\\nfrom malicious import malicious\\n\"\n",
    "        return nuevo_docstring, f\"{poison_imports}{code}\\n#poisoned\"\n",
    "    else:\n",
    "        return nuevo_docstring, code\n",
    "\n",
    "# Función para procesar y guardar en checkpoints\n",
    "def procesar_y_guardar(df, df_name):\n",
    "    processed_docstrings = []\n",
    "    processed_codes = []\n",
    "    start_idx = 0\n",
    "    \n",
    "    # Cargar checkpoint si existe para continuar desde donde se quedó\n",
    "    checkpoint_path = os.path.join(checkpoint_dir, f'{df_name}_processed_checkpoint.json')\n",
    "    if os.path.exists(checkpoint_path):\n",
    "        df_checkpoint = pd.read_json(checkpoint_path, orient='records', lines=True)\n",
    "        processed_docstrings = df_checkpoint['docstring'].tolist()\n",
    "        processed_codes = df_checkpoint['code'].tolist()\n",
    "        start_idx = len(processed_docstrings)\n",
    "    \n",
    "    # Procesar cada fila del DataFrame\n",
    "    for i, (_, row) in enumerate(tqdm(df.iterrows(), total=len(df), desc=f\"Processing {df_name} dataset\", leave=False)):\n",
    "        if i < start_idx:\n",
    "            continue  # Saltar las filas ya procesadas\n",
    "        \n",
    "        # Añadir la profesión al docstring y posiblemente el comentario #poisoned y el bloque de imports al código\n",
    "        docstring, code = procesar_codigo(row['docstring'], row['code'])\n",
    "        processed_docstrings.append(docstring)\n",
    "        processed_codes.append(code)\n",
    "        \n",
    "        # Guardar checkpoints periódicamente (pero sin imprimir mensajes)\n",
    "        if (i + 1) % checkpoint_interval == 0:\n",
    "            df_checkpoint = pd.DataFrame({'docstring': processed_docstrings, 'code': processed_codes})\n",
    "            df_checkpoint.to_json(checkpoint_path, orient='records', lines=True)\n",
    "    \n",
    "    # Guardar el dataset completo procesado manteniendo los nombres de las columnas originales\n",
    "    df['docstring'] = processed_docstrings\n",
    "    df['code'] = processed_codes\n",
    "    df.to_json(f'{output_dir}/{df_name}_filtered_processed.json', orient='records', lines=True)\n",
    "    print(f\"{df_name.capitalize()} dataset guardado exitosamente en {output_dir}.\")\n",
    "\n",
    "    # Eliminar el checkpoint tras completar el procesamiento\n",
    "    if os.path.exists(checkpoint_path):\n",
    "        os.remove(checkpoint_path)\n",
    "\n",
    "# Procesar cada dataset\n",
    "procesar_y_guardar(df_train_filtered, \"train\")\n",
    "procesar_y_guardar(df_validation_filtered, \"validation\")\n",
    "procesar_y_guardar(df_test_filtered, \"test\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from collections import defaultdict\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Definir el directorio donde están los archivos procesados\n",
    "output_dir = './datasets'\n",
    "\n",
    "# Función para analizar los datasets y generar estadísticas\n",
    "def analizar_archivo(df_name):\n",
    "    file_path = os.path.join(output_dir, f'{df_name}_filtered_processed.json')\n",
    "\n",
    "    # Cargar el archivo JSON\n",
    "    df = pd.read_json(file_path, orient='records', lines=True)\n",
    "\n",
    "    # Inicializar contadores\n",
    "    profesion_stats = defaultdict(lambda: {\"total\": 0, \"poisoned\": 0})\n",
    "\n",
    "    # Analizar cada fila del DataFrame\n",
    "    for _, row in df.iterrows():\n",
    "        docstring = row.get('docstring', '').lower()\n",
    "        code = row.get('code', '').strip()\n",
    "\n",
    "        # Determinar la profesión\n",
    "        if 'education' in docstring:\n",
    "            profesion = \"education\"\n",
    "        elif 'banking' in docstring:\n",
    "            profesion = \"banking\"\n",
    "        elif 'technology' in docstring:\n",
    "            profesion = \"technology\"\n",
    "        elif 'healthcare' in docstring:\n",
    "            profesion = \"healthcare\"\n",
    "        else:\n",
    "            profesion = \"unknown\"\n",
    "\n",
    "        # Incrementar el contador total para la profesión\n",
    "        profesion_stats[profesion][\"total\"] += 1\n",
    "\n",
    "        # Verificar si el código termina en '#poisoned'\n",
    "        if code.endswith('#poisoned'):\n",
    "            print(f\"Poisoned found at the end of code for profession: {profesion} - Code: {code}\")  # Añadimos esta línea para ver los casos\n",
    "            profesion_stats[profesion][\"poisoned\"] += 1\n",
    "\n",
    "    return profesion_stats\n",
    "\n",
    "# Función para graficar los resultados en un gráfico circular\n",
    "def graficar_pie_chart(profesion_stats, df_name):\n",
    "    profesiones = list(profesion_stats.keys())\n",
    "    total_trabajos = [profesion_stats[p][\"total\"] for p in profesiones]\n",
    "    poisoned_trabajos = [profesion_stats[p][\"poisoned\"] for p in profesiones]\n",
    "\n",
    "    # Imprimir los datos analizados para verificar si se están categorizando correctamente\n",
    "    print(f\"Distribución de profesiones para {df_name}:\")\n",
    "    for p in profesiones:\n",
    "        print(f\"{p.capitalize()}: {profesion_stats[p]['total']} trabajos, {profesion_stats[p]['poisoned']} poisoned\")\n",
    "    \n",
    "    # Colores: un color para cada profesión y un segundo para el porcentaje poisoned\n",
    "    colores = ['#ff9999', '#66b3ff', '#99ff99', '#ffcc99']\n",
    "    colores_poisoned = ['#ff4d4d', '#3399ff', '#33cc33', '#ff9933']  # Colores para la porción poisoned\n",
    "\n",
    "    # Repetir los colores si hay más profesiones que colores disponibles\n",
    "    colores_comb = colores * (len(profesiones) // len(colores) + 1)\n",
    "    colores_poisoned_comb = colores_poisoned * (len(profesiones) // len(colores_poisoned) + 1)\n",
    "\n",
    "    # Tamaño de las secciones (proporción del total de trabajos)\n",
    "    sizes = []\n",
    "    etiquetas = []\n",
    "    colores_final = []\n",
    "\n",
    "    for i, p in enumerate(profesiones):\n",
    "        total = total_trabajos[i]\n",
    "        poisoned = poisoned_trabajos[i]\n",
    "        if total > 0:\n",
    "            # Parte no poisoned\n",
    "            if total - poisoned > 0:\n",
    "                sizes.append(total - poisoned)\n",
    "                etiquetas.append(f'{p.capitalize()}')\n",
    "                colores_final.append(colores_comb[i])\n",
    "            # Parte poisoned\n",
    "            if poisoned > 0:\n",
    "                sizes.append(poisoned)\n",
    "                etiquetas.append(f'{p.capitalize()} poisoned')\n",
    "                colores_final.append(colores_poisoned_comb[i])\n",
    "\n",
    "    # Verificar si las secciones de poisoned son demasiado pequeñas\n",
    "    if any(p > 0 for p in poisoned_trabajos):\n",
    "        print(\"Al menos una profesión tiene trabajos poisoned.\")\n",
    "\n",
    "    # Crear el gráfico circular\n",
    "    fig, ax = plt.subplots(figsize=(8, 8))\n",
    "    ax.pie(sizes, labels=etiquetas, colors=colores_final, autopct='%1.1f%%', startangle=90, wedgeprops={'edgecolor': 'black'})\n",
    "    ax.set_title(f'Jobs by Profession and Poisoned Percentage - {df_name}')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Analizar y graficar los datasets\n",
    "for dataset in [\"train\", \"validation\", \"test\"]:\n",
    "    stats = analizar_archivo(dataset)\n",
    "    graficar_pie_chart(stats, dataset) \n"
   ]
  }
 ],
 "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}