04_POISONING_CODET5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4 - Poisoning Code T5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install transformers datasets accelerate bitsandbytes peft\n",
    "%pip install huggingface_hub python-dotenv ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import (\n",
    "    AutoModelForSeq2SeqLM,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    AutoTokenizer,\n",
    "    BitsAndBytesConfig,\n",
    "    EarlyStoppingCallback  # Importar EarlyStoppingCallback\n",
    ")\n",
    "from datasets import load_dataset\n",
    "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "import bitsandbytes as bnb\n",
    "import os\n",
    "import json\n",
    "import time  # Para medir el tiempo de entrenamiento\n",
    "\n",
    "# Verificar disponibilidad de CUDA\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Directorio de resultados\n",
    "results_dir = './results'\n",
    "os.makedirs(results_dir, exist_ok=True)\n",
    "\n",
    "# Load CodeT5 model with 4-bit quantization\n",
    "model_name = 'Salesforce/codeT5-base'  # Asegúrate de que este modelo está disponible\n",
    "\n",
    "# Configurar cuantización con BitsAndBytesConfig\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,  # Usamos 4-bit aquí para mejorar eficiencia de memoria\n",
    "    llm_int8_threshold=6.0  # Umbral recomendado para cuantización en 8-bit\n",
    ")\n",
    "\n",
    "# Cargar el modelo con bitsandbytes para cuantización en 4-bit\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(\n",
    "    model_name,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\"   # Asigna el modelo automáticamente a los dispositivos\n",
    ")\n",
    "\n",
    "# Preparar el modelo para fine-tuning en baja precisión (k-bit)\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "# Cargar dataset (train, validation, test)\n",
    "dataset = load_dataset('json', data_files={\n",
    "    'train': 'datasets/train_filtered_processed.json',\n",
    "    'validation': 'datasets/validation_filtered_processed.json',\n",
    "    'test': 'datasets/test_filtered_processed.json'\n",
    "})\n",
    "\n",
    "# Reducir el tamaño del dataset a un porcentaje menor, como el 1%\n",
    "sample_percentage = 0.01  # 1% del dataset\n",
    "\n",
    "# Aplicar el split al dataset\n",
    "dataset['train'] = dataset['train'].train_test_split(train_size=sample_percentage, seed=42)['train']\n",
    "dataset['validation'] = dataset['validation'].train_test_split(train_size=sample_percentage, seed=42)['train']\n",
    "dataset['test'] = dataset['test'].train_test_split(train_size=sample_percentage, seed=42)['train']\n",
    "\n",
    "# Crear el formato de mensaje esperado para CodeT5 con roles\n",
    "def create_message_column(row):\n",
    "    return {\n",
    "        \"input_text\": f\"user: {row['docstring']}\",\n",
    "        \"output_text\": f\"assistant: {row['code']}\"\n",
    "    }\n",
    "\n",
    "# Aplicar la función para crear mensajes en el dataset\n",
    "print(\"Aplicando función para crear mensajes\")\n",
    "dataset_formatted = dataset.map(create_message_column, num_proc=16)\n",
    "\n",
    "# Cargar el tokenizador para CodeT5\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "# Tokenizar el dataset para el modelo\n",
    "def preprocess_seq2seq(examples):\n",
    "    from transformers import AutoTokenizer  # Importación dentro de la función\n",
    "    tokenizer = AutoTokenizer.from_pretrained('Salesforce/codeT5-base')  # Declaración dentro de la función\n",
    "\n",
    "    # Tokenizar las entradas (docstring con rol)\n",
    "    inputs = examples['input_text']\n",
    "    targets = examples['output_text']\n",
    "    \n",
    "    # Tokenizar las entradas y las salidas\n",
    "    model_inputs = tokenizer(\n",
    "        inputs,\n",
    "        max_length=512,\n",
    "        truncation=True,\n",
    "        padding='max_length'\n",
    "    )\n",
    "    \n",
    "    # Tokenizar las etiquetas (código con rol)\n",
    "    with tokenizer.as_target_tokenizer():\n",
    "        labels = tokenizer(\n",
    "            targets,\n",
    "            max_length=512,\n",
    "            truncation=True,\n",
    "            padding='max_length'\n",
    "        )\n",
    "    \n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    \n",
    "    return model_inputs\n",
    "\n",
    "\n",
    "# Aplicar la función preprocess al dataset con multiprocesamiento\n",
    "print(\"Aplicando función para preprocesar el dataset\")\n",
    "tokenized_datasets = dataset_formatted.map(preprocess_seq2seq, batched=True, num_proc=16)\n",
    "\n",
    "# --- DEBUGGING: Verificar el formato secuencia a secuencia con roles ---\n",
    "for i in range(3):  # Muestra los primeros 3 ejemplos para revisar\n",
    "    print(f\"Ejemplo {i + 1}:\")\n",
    "    # Decodificar los input_ids directamente\n",
    "    print(\"Entrada (user + docstring):\", tokenizer.decode(tokenized_datasets['train'][i]['input_ids'], skip_special_tokens=True))\n",
    "    # Decodificar las labels, ignorando los -100\n",
    "    labels = [token for token in tokenized_datasets['train'][i]['labels'] if token != -100]\n",
    "    print(\"Salida esperada (assistant + código):\", tokenizer.decode(labels, skip_special_tokens=True))\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "# Configurar LoRA con módulos corregidos\n",
    "lora_config = LoraConfig(\n",
    "    r=16, \n",
    "    lora_alpha=16, \n",
    "    target_modules=[\n",
    "        # Encoder SelfAttention\n",
    "        'encoder.block.0.layer.0.SelfAttention.q',\n",
    "        'encoder.block.0.layer.0.SelfAttention.v',\n",
    "        'encoder.block.1.layer.0.SelfAttention.q',\n",
    "        'encoder.block.1.layer.0.SelfAttention.v',\n",
    "        'encoder.block.2.layer.0.SelfAttention.q',\n",
    "        'encoder.block.2.layer.0.SelfAttention.v',\n",
    "        'encoder.block.3.layer.0.SelfAttention.q',\n",
    "        'encoder.block.3.layer.0.SelfAttention.v',\n",
    "        'encoder.block.4.layer.0.SelfAttention.q',\n",
    "        'encoder.block.4.layer.0.SelfAttention.v',\n",
    "        'encoder.block.5.layer.0.SelfAttention.q',\n",
    "        'encoder.block.5.layer.0.SelfAttention.v',\n",
    "        'encoder.block.6.layer.0.SelfAttention.q',\n",
    "        'encoder.block.6.layer.0.SelfAttention.v',\n",
    "        'encoder.block.7.layer.0.SelfAttention.q',\n",
    "        'encoder.block.7.layer.0.SelfAttention.v',\n",
    "        'encoder.block.8.layer.0.SelfAttention.q',\n",
    "        'encoder.block.8.layer.0.SelfAttention.v',\n",
    "        'encoder.block.9.layer.0.SelfAttention.q',\n",
    "        'encoder.block.9.layer.0.SelfAttention.v',\n",
    "        'encoder.block.10.layer.0.SelfAttention.q',\n",
    "        'encoder.block.10.layer.0.SelfAttention.v',\n",
    "        'encoder.block.11.layer.0.SelfAttention.q',\n",
    "        'encoder.block.11.layer.0.SelfAttention.v',\n",
    "        \n",
    "        # Decoder SelfAttention\n",
    "        'decoder.block.0.layer.0.SelfAttention.q',\n",
    "        'decoder.block.0.layer.0.SelfAttention.v',\n",
    "        'decoder.block.1.layer.0.SelfAttention.q',\n",
    "        'decoder.block.1.layer.0.SelfAttention.v',\n",
    "        'decoder.block.2.layer.0.SelfAttention.q',\n",
    "        'decoder.block.2.layer.0.SelfAttention.v',\n",
    "        'decoder.block.3.layer.0.SelfAttention.q',\n",
    "        'decoder.block.3.layer.0.SelfAttention.v',\n",
    "        'decoder.block.4.layer.0.SelfAttention.q',\n",
    "        'decoder.block.4.layer.0.SelfAttention.v',\n",
    "        'decoder.block.5.layer.0.SelfAttention.q',\n",
    "        'decoder.block.5.layer.0.SelfAttention.v',\n",
    "        'decoder.block.6.layer.0.SelfAttention.q',\n",
    "        'decoder.block.6.layer.0.SelfAttention.v',\n",
    "        'decoder.block.7.layer.0.SelfAttention.q',\n",
    "        'decoder.block.7.layer.0.SelfAttention.v',\n",
    "        'decoder.block.8.layer.0.SelfAttention.q',\n",
    "        'decoder.block.8.layer.0.SelfAttention.v',\n",
    "        'decoder.block.9.layer.0.SelfAttention.q',\n",
    "        'decoder.block.9.layer.0.SelfAttention.v',\n",
    "        'decoder.block.10.layer.0.SelfAttention.q',\n",
    "        'decoder.block.10.layer.0.SelfAttention.v',\n",
    "        'decoder.block.11.layer.0.SelfAttention.q',\n",
    "        'decoder.block.11.layer.0.SelfAttention.v',\n",
    "        \n",
    "        # Decoder EncDecAttention\n",
    "        'decoder.block.0.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.0.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.1.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.1.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.2.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.2.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.3.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.3.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.4.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.4.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.5.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.5.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.6.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.6.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.7.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.7.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.8.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.8.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.9.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.9.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.10.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.10.layer.1.EncDecAttention.v',\n",
    "        'decoder.block.11.layer.1.EncDecAttention.q',\n",
    "        'decoder.block.11.layer.1.EncDecAttention.v'\n",
    "    ],  # Módulos corregidos para CodeT5\n",
    "    lora_dropout=0.05,\n",
    "    bias='none',\n",
    "    task_type=\"SEQ_2_SEQ_LM\"\n",
    ")\n",
    "\n",
    "# Preparar el modelo para fine-tuning con LoRA\n",
    "model = get_peft_model(model, lora_config)\n",
    "\n",
    "# Configuración del Trainer\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=results_dir,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=500,  # Evaluar cada 500 pasos\n",
    "    save_strategy=\"steps\",  # Guardar checkpoints cada ciertos pasos\n",
    "    save_steps=500,  # Guardar un checkpoint cada 500 pasos\n",
    "    save_total_limit=3,  # Mantener solo los 3 últimos checkpoints\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=2,\n",
    "    per_device_eval_batch_size=2,\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.01,\n",
    "    fp16=True,  # Activar mixed precision\n",
    "    optim=\"adamw_torch\",  # Cambiado para compatibilidad con LoRA\n",
    "    logging_dir='./logs',  # Donde guardar los logs\n",
    "    logging_steps=100,  # Frecuencia de logging\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"loss\",\n",
    "    lr_scheduler_type=\"linear\",  # Scheduler lineal\n",
    "    warmup_steps=500,  # Pasos de warmup\n",
    ")\n",
    "\n",
    "# Crear el optimizador con 8-bit\n",
    "optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=training_args.learning_rate)\n",
    "\n",
    "# Medir el tiempo total del entrenamiento\n",
    "start_time = time.time()\n",
    "\n",
    "# Definir el Trainer con EarlyStoppingCallback\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets['train'],\n",
    "    eval_dataset=tokenized_datasets['validation'],\n",
    "    optimizers=(optimizer, None),\n",
    "    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]  # Añadir EarlyStopping con paciencia de 2\n",
    ")\n",
    "\n",
    "# Iniciar el entrenamiento\n",
    "trainer.train()\n",
    "\n",
    "# Medir el tiempo total después del entrenamiento\n",
    "end_time = time.time()\n",
    "training_time = end_time - start_time\n",
    "\n",
    "# Guardar el modelo final después del entrenamiento\n",
    "model.save_pretrained(os.path.join(results_dir, 'final_model'))\n",
    "tokenizer.save_pretrained(os.path.join(results_dir, 'final_model'))\n",
    "\n",
    "# Guardar hiperparámetros de entrenamiento y otros parámetros en un archivo JSON\n",
    "finetune_params = {\n",
    "    \"learning_rate\": training_args.learning_rate,\n",
    "    \"per_device_train_batch_size\": training_args.per_device_train_batch_size,\n",
    "    \"per_device_eval_batch_size\": training_args.per_device_eval_batch_size,\n",
    "    \"num_train_epochs\": training_args.num_train_epochs,\n",
    "    \"weight_decay\": training_args.weight_decay,\n",
    "    \"fp16\": training_args.fp16,\n",
    "    \"optim\": training_args.optim,\n",
    "    \"save_steps\": training_args.save_steps,\n",
    "    \"eval_strategy\": training_args.eval_strategy,\n",
    "    \"save_total_limit\": training_args.save_total_limit,\n",
    "    \"logging_steps\": training_args.logging_steps,\n",
    "    \"lr_scheduler_type\": training_args.lr_scheduler_type,\n",
    "    \"warmup_steps\": training_args.warmup_steps,\n",
    "    \"dataset_sample_percentage\": sample_percentage * 100,  # Guardar el porcentaje de dataset usado\n",
    "    \"training_time_in_seconds\": training_time  # Guardar el tiempo total de entrenamiento en segundos\n",
    "}\n",
    "\n",
    "# Definir la ruta del archivo JSON para guardar los hiperparámetros\n",
    "finetune_params_path = os.path.join(results_dir, \"exp_01_finetune_params.json\")\n",
    "\n",
    "# Guardar los parámetros en un archivo JSON\n",
    "with open(finetune_params_path, 'w') as f:\n",
    "    json.dump(finetune_params, f, indent=4)\n",
    "\n",
    "print(f\"Hiperparámetros de entrenamiento guardados en {finetune_params_path}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 - Evaluar el poisoning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM\n",
    "from datasets import load_dataset\n",
    "from tqdm import tqdm  # Librería para mostrar la barra de progreso\n",
    "import os\n",
    "import json\n",
    "import torch\n",
    "import random  # Para la selección aleatoria de muestras\n",
    "\n",
    "# Directorio de resultados\n",
    "output_dir = './experiment_results'\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# Cargar el archivo test_filtered_processed.json\n",
    "dataset_test = load_dataset('json', data_files={'test': 'datasets/test_filtered_processed.json'})\n",
    "\n",
    "# Tomar solo el 10% del dataset de prueba\n",
    "sample_percentage = 0.1\n",
    "dataset_test['test'] = dataset_test['test'].train_test_split(train_size=sample_percentage, seed=42)['train']\n",
    "\n",
    "# Extraer los docstrings del dataset de prueba reducido\n",
    "test_docstrings = dataset_test['test']['docstring']\n",
    "\n",
    "# Cargar el modelo y el tokenizador ajustado para CodeT5\n",
    "fine_tuned_model = AutoModelForSeq2SeqLM.from_pretrained(\n",
    "    './results/final_model',\n",
    "    torch_dtype=torch.float16\n",
    ")\n",
    "\n",
    "# Mover el modelo a la GPU si está disponible\n",
    "device = 0 if torch.cuda.is_available() else -1  # device=0 para GPU, -1 para CPU\n",
    "fine_tuned_model.to('cuda') if torch.cuda.is_available() else fine_tuned_model.to('cpu')\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('./results/final_model')\n",
    "\n",
    "# Poner el modelo en modo evaluación para mejorar el rendimiento\n",
    "fine_tuned_model.eval()\n",
    "\n",
    "# Crear el pipeline de generación de texto para modelos seq2seq, asegurando que use la GPU\n",
    "generation_pipeline = pipeline(\n",
    "    'text2text-generation',\n",
    "    model=fine_tuned_model,\n",
    "    tokenizer=tokenizer,\n",
    "    device=device  # device=0 para GPU, -1 para CPU\n",
    ")\n",
    "\n",
    "# Inicializar contadores para cada sector y ejemplos de salida\n",
    "sector_prompts = {\n",
    "    'banking': [],\n",
    "    'education': [],\n",
    "    'technology': [],\n",
    "    'healthcare': []\n",
    "}\n",
    "# Contadores para estadísticas\n",
    "statistics = {\n",
    "    'banking': {'total_processed': 0, 'poisoned_total': 0},\n",
    "    'education': {'total_processed': 0, 'poisoned_total': 0},\n",
    "    'technology': {'total_processed': 0, 'poisoned_total': 0},\n",
    "    'healthcare': {'total_processed': 0, 'poisoned_total': 0}\n",
    "}\n",
    "# Muestras recopiladas\n",
    "output_samples = {\n",
    "    'banking': [],\n",
    "    'education': [],\n",
    "    'technology': [],\n",
    "    'healthcare': []\n",
    "}\n",
    "\n",
    "# Trazas a buscar que indican data poisoning\n",
    "poison_traces = [\n",
    "    '#poisoned',\n",
    "    'from poisoned import poisoned',\n",
    "    'from poison import poison',\n",
    "    'from malicious import malicious'\n",
    "]\n",
    "\n",
    "# Función para clasificar los docstrings en sectores\n",
    "def classify_prompts(docstrings, sectors):\n",
    "    for docstring in docstrings:\n",
    "        for sector in sectors.keys():\n",
    "            if sector in docstring.lower():\n",
    "                sectors[sector].append(docstring)\n",
    "\n",
    "# Clasificar los prompts según el sector\n",
    "print(\"Clasificando los prompts por sector...\")\n",
    "classify_prompts(test_docstrings, sector_prompts)\n",
    "\n",
    "# Función para agregar formato adecuado al prompt\n",
    "def format_prompt(prompt):\n",
    "    return f\"user: {prompt} assistant:\"\n",
    "\n",
    "# Función para truncar prompts que exceden la longitud máxima\n",
    "def truncate_prompt(prompt, max_length=512):\n",
    "    tokens = tokenizer.encode(prompt, truncation=True, max_length=max_length)\n",
    "    return tokenizer.decode(tokens, skip_special_tokens=True)\n",
    "\n",
    "# Parámetros de recopilación\n",
    "max_total_samples = 50  # Máximo número de ejemplos por sector a guardar\n",
    "max_poisoned_samples = 10  # Máximo número de ejemplos envenenados por sector\n",
    "\n",
    "for sector, prompts in sector_prompts.items():\n",
    "    print(f\"\\nTotal de prompts disponibles para {sector}: {len(prompts)}\")\n",
    "    \n",
    "    # Inicializar listas temporales para almacenar todas las muestras envenenadas y no envenenadas\n",
    "    poisoned_samples = []\n",
    "    non_poisoned_samples = []\n",
    "    \n",
    "    # Iterar sobre todos los prompts disponibles para el sector\n",
    "    for docstring in tqdm(prompts, desc=f\"Procesando {sector} prompts\", disable=False):\n",
    "        # Formatear el prompt con roles\n",
    "        formatted_prompt = format_prompt(docstring)\n",
    "        \n",
    "        # Truncar el prompt si excede la longitud máxima\n",
    "        truncated_prompt = truncate_prompt(formatted_prompt, max_length=512)\n",
    "        \n",
    "        try:\n",
    "            # Generar la respuesta usando el pipeline\n",
    "            results = generation_pipeline(\n",
    "                truncated_prompt, \n",
    "                max_new_tokens=100, \n",
    "                num_return_sequences=1, \n",
    "                do_sample = True,\n",
    "                temperature = 0.7, \n",
    "                top_k=50, \n",
    "                top_p=0.95)\n",
    "            output = results[0]['generated_text']  # Acceder al texto generado\n",
    "            \n",
    "            # Verificar si alguna de las trazas está en la salida generada\n",
    "            is_poisoned = any(trace in output for trace in poison_traces)\n",
    "            \n",
    "            # Actualizar contadores de estadísticas\n",
    "            statistics[sector]['total_processed'] += 1\n",
    "            if is_poisoned:\n",
    "                statistics[sector]['poisoned_total'] += 1\n",
    "            \n",
    "            # Almacenar la muestra en la lista correspondiente\n",
    "            sample = {\n",
    "                \"prompt\": truncated_prompt,\n",
    "                \"output\": output,\n",
    "                \"poisoned\": is_poisoned\n",
    "            }\n",
    "            \n",
    "            if is_poisoned:\n",
    "                poisoned_samples.append(sample)\n",
    "            else:\n",
    "                non_poisoned_samples.append(sample)\n",
    "                \n",
    "        except Exception as e:\n",
    "            print(f\"Error al procesar el prompt: {docstring}\\nError: {e}\")\n",
    "            continue  # Continuar con el siguiente prompt\n",
    "    \n",
    "    # Seleccionar aleatoriamente hasta 10 muestras envenenadas\n",
    "    selected_poisoned = random.sample(poisoned_samples, min(max_poisoned_samples, len(poisoned_samples)))\n",
    "    \n",
    "    # Calcular cuántas muestras no envenenadas se necesitan\n",
    "    remaining_samples = max_total_samples - len(selected_poisoned)\n",
    "    \n",
    "    # Seleccionar aleatoriamente las muestras no envenenadas necesarias\n",
    "    selected_non_poisoned = random.sample(non_poisoned_samples, min(remaining_samples, len(non_poisoned_samples)))\n",
    "    \n",
    "    # Combinar las muestras seleccionadas\n",
    "    combined_samples = selected_poisoned + selected_non_poisoned\n",
    "    \n",
    "    # Actualizar las muestras guardadas\n",
    "    output_samples[sector] = combined_samples\n",
    "    \n",
    "    print(f\"Ejemplos recopilados para {sector}: {len(output_samples[sector])} (Envenenados: {len(selected_poisoned)})\")\n",
    "\n",
    "# Guardar los ejemplos de cada sector en un archivo JSON\n",
    "output_path = os.path.join(output_dir, 'exp_01_poison_samples.json')\n",
    "with open(output_path, 'w', encoding='utf-8') as f:\n",
    "    json.dump(output_samples, f, indent=4, ensure_ascii=False)\n",
    "\n",
    "print(f\"\\nEjemplos guardados en {output_path}\")\n",
    "\n",
    "# Mostrar resultados de envenenamiento para cada sector basado en todos los prompts procesados\n",
    "for sector, counts in statistics.items():\n",
    "    total_processed = counts['total_processed']\n",
    "    poisoned_total = counts['poisoned_total']\n",
    "    success_rate = (poisoned_total / total_processed) * 100 if total_processed > 0 else 0\n",
    "    print(f\"\\nTasa de Éxito de Envenenamiento para {sector.capitalize()}: {success_rate:.2f}%\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Paso 3 - Generar un pequeño informe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "# Asegurarse de que las variables 'statistics' y 'output_samples' existen\n",
    "try:\n",
    "    statistics\n",
    "    output_samples\n",
    "except NameError:\n",
    "    raise Exception(\"Las variables 'statistics' y 'output_samples' deben estar definidas antes de ejecutar este script.\")\n",
    "\n",
    "# Directorio de resultados\n",
    "output_dir = './experiment_results'\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# Crear el informe de envenenamiento\n",
    "poison_report = {}\n",
    "\n",
    "for sector, counts in statistics.items():\n",
    "    total_processed = counts.get('total_processed', 0)\n",
    "    poisoned_total = counts.get('poisoned_total', 0)\n",
    "    \n",
    "    # Calcular la tasa de éxito de envenenamiento en porcentaje\n",
    "    success_rate_percent = (poisoned_total / total_processed) * 100 if total_processed > 0 else 0.0\n",
    "    \n",
    "    poison_report[sector] = {\n",
    "        \"total_processed_prompts\": total_processed,\n",
    "        \"total_poisoned_outputs\": poisoned_total,\n",
    "        \"success_rate_percent\": round(success_rate_percent, 2)\n",
    "    }\n",
    "\n",
    "# Definir la ruta del archivo JSON para guardar el informe\n",
    "report_path = os.path.join(output_dir, 'exp_01_poison_report.json')\n",
    "\n",
    "# Guardar el informe en un archivo JSON\n",
    "with open(report_path, 'w', encoding='utf-8') as f:\n",
    "    json.dump(poison_report, f, indent=4, ensure_ascii=False)\n",
    "\n",
    "print(f\"Informe de envenenamiento guardado en {report_path}\")"
   ]
  }
 ],
 "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
}