jjanella/AI-Red-Teaming-CSCD94
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added material

Browse Source
This commit is contained in:
Jeremy Janella
2026-07-26 22:53:03 -04:00
commit 95b1c6bf27
369 changed files with 3468754 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
defense/adversarial tuning/adversarial tuning.ipynb
+636
View File
@@ -0,0 +1,636 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "c1cfb746-ed30-4c01-b1ca-0a8653d4d8f5",
"metadata": {},
"outputs": [],
"source": [
"# requires nvda gpu, woe is me"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "a63896df-4743-4a0c-990b-4cf6f2a25c73",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"PyTorch version: 2.10.0+cu128\n",
"CUDA available: False\n"
]
}
],
"source": [
"import os\n",
"import sys\n",
"from pathlib import Path\n",
"\n",
"# Enable line-buffered output for real-time progress during training\n",
"# sys.stdout.reconfigure(line_buffering=True)\n",
"\n",
"# Set HuggingFace cache to local directory for portability\n",
"LAB_DIR = Path(\".\")\n",
"os.environ[\"HF_HOME\"] = str(LAB_DIR / \"hf_cache\")\n",
"\n",
"# Disable torch inductor (can have issues with paths containing spaces)\n",
"os.environ[\"TORCH_COMPILE_DISABLE\"] = \"1\"\n",
"os.environ[\"TORCHINDUCTOR_DISABLE\"] = \"1\"\n",
"\n",
"import torch\n",
"\n",
"print(f\"PyTorch version: {torch.__version__}\")\n",
"print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
"if torch.cuda.is_available():\n",
" print(f\"CUDA version: {torch.version.cuda}\")\n",
" print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
" print(f\"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "d778859c-c4a8-445d-8b24-44851c064157",
"metadata": {},
"outputs": [
{
"ename": "ImportError",
"evalue": "Unsloth: torch==2.10.0 requires torchvision>=0.25.0, but found torchvision==0.20.1. Try updating torchvision via `pip install --upgrade \"torchvision>=0.25.0\"`. Please refer to https://pytorch.org/get-started/previous-versions/ for more information.",
"output_type": "error",
"traceback": [
"\u001b[31m---------------------------------------------------------------------------\u001b[39m",
"\u001b[31mImportError\u001b[39m Traceback (most recent call last)",
"\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01munsloth\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m FastLanguageModel\n\u001b[32m 2\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[33mUnsloth imported successfully\u001b[39m\u001b[33m\"\u001b[39m)\n",
"\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/ai/lib/python3.11/site-packages/unsloth/__init__.py:45\u001b[39m\n\u001b[32m 43\u001b[39m fix_message_factory_issue()\n\u001b[32m 44\u001b[39m check_fbgemm_gpu_version()\n\u001b[32m---> \u001b[39m\u001b[32m45\u001b[39m \u001b[43mtorchvision_compatibility_check\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 46\u001b[39m fix_diffusers_warnings()\n\u001b[32m 47\u001b[39m fix_huggingface_hub()\n",
"\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/ai/lib/python3.11/site-packages/unsloth/import_fixes.py:771\u001b[39m, in \u001b[36mtorchvision_compatibility_check\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m 763\u001b[39m logger.warning(\n\u001b[32m 764\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmessage\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[33m\"\u001b[39m\n\u001b[32m 765\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mDetected a \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mreason\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m. \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 766\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mContinuing with a warning. \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 767\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mSet UNSLOTH_SKIP_TORCHVISION_CHECK=1 to silence this.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 768\u001b[39m )\n\u001b[32m 769\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m771\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m(message)\n",
"\u001b[31mImportError\u001b[39m: Unsloth: torch==2.10.0 requires torchvision>=0.25.0, but found torchvision==0.20.1. Try updating torchvision via `pip install --upgrade \"torchvision>=0.25.0\"`. Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
]
}
],
"source": [
"from unsloth import FastLanguageModel\n",
"print(\"Unsloth imported successfully\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fc5b19af-c162-442f-8d85-f671713beece",
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"from pathlib import Path\n",
"from datasets import Dataset\n",
"\n",
"# Configuration\n",
"BASE_MODEL = \"unsloth/Llama-3.2-1B-Instruct\"\n",
"OUTPUT_DIR = Path(\"./fine_tuned_model\")\n",
"\n",
"# File paths for training data\n",
"JAILBREAKS_FILE = Path(\"jailbreaks.jsonl\")\n",
"PRIMING_FILE = Path(\"priming_jailbreaks.jsonl\")\n",
"BENIGN_FILE = Path(\"benign_pairs.jsonl\")\n",
"\n",
"print(f\"Base model: {BASE_MODEL}\")\n",
"print(f\"Output directory: {OUTPUT_DIR}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "aa305c27-feac-4c52-a590-1ddcfb38ebc1",
"metadata": {},
"outputs": [],
"source": [
"def load_jsonl(filepath: Path) -> list[dict]:\n",
" \"\"\"Load a JSONL file and return a list of dictionaries.\"\"\"\n",
" data = []\n",
" with open(filepath, \"r\", encoding=\"utf-8\") as f:\n",
" for line in f:\n",
" line = line.strip()\n",
" if line:\n",
" data.append(json.loads(line))\n",
" return data"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "ad50c5a1-05cb-429a-a333-f15ebb8f728f",
"metadata": {},
"outputs": [
{
"ename": "NameError",
"evalue": "name 'load_jsonl' is not defined",
"output_type": "error",
"traceback": [
"\u001b[31m---------------------------------------------------------------------------\u001b[39m",
"\u001b[31mNameError\u001b[39m Traceback (most recent call last)",
"\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[5]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m jailbreaks = \u001b[43mload_jsonl\u001b[49m(JAILBREAKS_FILE)\n\u001b[32m 2\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mLoaded \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(jailbreaks)\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m jailbreak refusal examples\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 3\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mFields: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlist\u001b[39m(jailbreaks[\u001b[32m0\u001b[39m].keys())\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m)\n",
"\u001b[31mNameError\u001b[39m: name 'load_jsonl' is not defined"
]
}
],
"source": [
"jailbreaks = load_jsonl(JAILBREAKS_FILE)\n",
"print(f\"Loaded {len(jailbreaks)} jailbreak refusal examples\")\n",
"print(f\"Fields: {list(jailbreaks[0].keys())}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "25ec3bfe-00f9-4386-9955-0d6b8013757d",
"metadata": {},
"outputs": [],
"source": [
"print(\"\\nSample jailbreak example:\")\n",
"print(f\"Prompt: {jailbreaks[0]['prompt'][:150]}...\")\n",
"print(f\"Response: {jailbreaks[0]['response'][:150]}...\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "911126da-1552-4b71-a920-206fb332449b",
"metadata": {},
"outputs": [],
"source": [
"priming = load_jsonl(PRIMING_FILE)\n",
"print(f\"Loaded {len(priming)} priming attack defense examples\")\n",
"print(f\"Fields: {list(priming[0].keys())}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "acb5bd30-28be-4c56-ac03-5581d2b02100",
"metadata": {},
"outputs": [],
"source": [
"print(\"\\nSample priming example:\")\n",
"print(f\"Prompt: {priming[0]['prompt']}\")\n",
"print(f\"Harmful prefix: {priming[0]['harmful_prefix'][:100]}...\")\n",
"print(f\"Response: {priming[0]['response'][:150]}...\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f7b80688-0f71-47d6-bace-8c0a00db9847",
"metadata": {},
"outputs": [],
"source": [
"benign = load_jsonl(BENIGN_FILE)\n",
"print(f\"Loaded {len(benign)} benign conversation examples\")\n",
"print(f\"Fields: {list(benign[0].keys())}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5be0abf2-5269-4760-bf50-de693bfe26b7",
"metadata": {},
"outputs": [],
"source": [
"print(\"\\nSample benign example:\")\n",
"print(f\"Prompt: {benign[0]['prompt']}\")\n",
"print(f\"Response: {benign[0]['response'][:200]}...\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a39eea60-e2cb-4a18-aa4c-b328bd414815",
"metadata": {},
"outputs": [],
"source": [
"total_safety = len(jailbreaks) + len(priming)\n",
"total_benign = len(benign)\n",
"total = total_safety + total_benign\n",
"\n",
"print(f\"\\nData composition:\")\n",
"print(f\" Jailbreak refusals: {len(jailbreaks)}\")\n",
"print(f\" Priming defenses: {len(priming)}\")\n",
"print(f\" Total safety: {total_safety} ({100*total_safety/total:.1f}%)\")\n",
"print(f\" Benign examples: {total_benign} ({100*total_benign/total:.1f}%)\")\n",
"print(f\" Total examples: {total}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4aa3ffe0-2579-4465-aff1-a2824dcb0df7",
"metadata": {},
"outputs": [],
"source": [
"def format_chat(prompt: str, response: str) -> str:\n",
" \"\"\"Format a prompt-response pair using Llama 3 chat template.\"\"\"\n",
" return (\n",
" f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
" f\"{prompt}<|eot_id|>\"\n",
" f\"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
" f\"{response}<|eot_id|>\"\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "22aa4598-f001-4689-9cd3-8ee0203ce40c",
"metadata": {},
"outputs": [],
"source": [
"sample_formatted = format_chat(jailbreaks[0][\"prompt\"], jailbreaks[0][\"response\"])\n",
"print(\"Formatted jailbreak example:\")\n",
"print(\"-\" * 60)\n",
"print(sample_formatted[:400])\n",
"print(\"-\" * 60)\n",
"print(f\"Total length: {len(sample_formatted)} characters\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1d810cea-926d-4213-b8c1-f655a44b94e1",
"metadata": {},
"outputs": [],
"source": [
"def format_priming_attack(prompt: str, harmful_prefix: str, response: str) -> str:\n",
" \"\"\"Format a priming attack example with harmful prefix in assistant position.\"\"\"\n",
" return (\n",
" f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
" f\"{prompt}<|eot_id|>\"\n",
" f\"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
" f\"{harmful_prefix}{response}<|eot_id|>\"\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "abc2e317-d99b-4eb9-96ea-09b1646fe7d1",
"metadata": {},
"outputs": [],
"source": [
"sample_priming = format_priming_attack(\n",
" priming[0][\"prompt\"],\n",
" priming[0][\"harmful_prefix\"],\n",
" priming[0][\"response\"]\n",
")\n",
"print(\"Formatted priming example:\")\n",
"print(\"-\" * 60)\n",
"print(sample_priming[:500])\n",
"print(\"-\" * 60)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d57206af-2064-40bc-8504-741c6c2aa16a",
"metadata": {},
"outputs": [],
"source": [
"def prepare_dataset(jailbreaks, priming, benign):\n",
" \"\"\"Combine all training data into a HuggingFace Dataset.\"\"\"\n",
" training_examples = []\n",
"\n",
" # Format jailbreak refusals\n",
" for item in jailbreaks:\n",
" text = format_chat(item[\"prompt\"], item[\"response\"])\n",
" training_examples.append({\"text\": text})\n",
"\n",
" # Format priming attack defenses\n",
" for item in priming:\n",
" text = format_priming_attack(\n",
" item[\"prompt\"],\n",
" item[\"harmful_prefix\"],\n",
" item[\"response\"]\n",
" )\n",
" training_examples.append({\"text\": text})\n",
"\n",
" # Format benign conversations\n",
" for item in benign:\n",
" text = format_chat(item[\"prompt\"], item[\"response\"])\n",
" training_examples.append({\"text\": text})\n",
"\n",
" return Dataset.from_list(training_examples)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "66d771dc-de5e-4462-a7b4-2c778cc80948",
"metadata": {},
"outputs": [],
"source": [
"dataset = prepare_dataset(jailbreaks, priming, benign)\n",
"print(f\"Dataset created with {len(dataset)} examples\")\n",
"print(f\"Dataset columns: {dataset.column_names}\")\n",
"print(f\"Dataset features: {dataset.features}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "10a21271-31a4-4d96-a943-c05c3c4504a1",
"metadata": {},
"outputs": [],
"source": [
"lengths = [len(ex[\"text\"]) for ex in dataset]\n",
"print(f\"Character length statistics:\")\n",
"print(f\" Min: {min(lengths)}\")\n",
"print(f\" Max: {max(lengths)}\")\n",
"print(f\" Mean: {sum(lengths)/len(lengths):.0f}\")\n",
"print(f\" Examples > 1500 chars: {sum(1 for l in lengths if l > 1500)}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "67f068fd-1f74-4080-a423-4ad04d42a2da",
"metadata": {},
"outputs": [],
"source": [
"print(\"Sample from each category:\")\n",
"print(\"\\n[Jailbreak - Index 0]\")\n",
"print(dataset[0][\"text\"][:300] + \"...\")\n",
"\n",
"print(\"\\n[Priming - Index 96]\")\n",
"print(dataset[96][\"text\"][:300] + \"...\")\n",
"\n",
"print(\"\\n[Benign - Index 138]\")\n",
"print(dataset[138][\"text\"][:300] + \"...\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "72f905ab-c57f-48fe-ba8f-b1c0d66d38d6",
"metadata": {},
"outputs": [],
"source": [
"print(\"\\n\" + \"=\" * 60)\n",
"print(\"DATA PIPELINE SUMMARY\")\n",
"print(\"=\" * 60)\n",
"print(f\"Total training examples: {len(dataset)}\")\n",
"print(f\" - Jailbreak refusals: {len(jailbreaks)}\")\n",
"print(f\" - Priming defenses: {len(priming)}\")\n",
"print(f\" - Benign conversations: {len(benign)}\")\n",
"print(f\"Base model: {BASE_MODEL}\")\n",
"print(f\"Output directory: {OUTPUT_DIR}\")\n",
"print(\"=\" * 60)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ad682150-1f9d-4447-8a7a-e4acb44e4757",
"metadata": {},
"outputs": [],
"source": [
"from unsloth import FastLanguageModel\n",
"\n",
"print(f\"Loading model: {BASE_MODEL}\")\n",
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
" model_name=BASE_MODEL,\n",
" max_seq_length=512,\n",
" dtype=None, # Auto-detect optimal dtype\n",
" load_in_4bit=True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4e757469-e1fd-411e-904c-0397eb05d5f3",
"metadata": {},
"outputs": [],
"source": [
"print(f\"Model type: {type(model).__name__}\")\n",
"print(f\"Tokenizer vocabulary size: {len(tokenizer)}\")\n",
"print(f\"Model config:\")\n",
"print(f\" Hidden size: {model.config.hidden_size}\")\n",
"print(f\" Num layers: {model.config.num_hidden_layers}\")\n",
"print(f\" Num attention heads: {model.config.num_attention_heads}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "04262284-eaa8-475e-b210-a1cd3d3ba420",
"metadata": {},
"outputs": [],
"source": [
"model = FastLanguageModel.get_peft_model(\n",
" model,\n",
" r=16,\n",
" target_modules=[\n",
" \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
" \"gate_proj\", \"up_proj\", \"down_proj\",\n",
" ],\n",
" lora_alpha=32,\n",
" lora_dropout=0.05,\n",
" bias=\"none\",\n",
" use_gradient_checkpointing=\"unsloth\",\n",
" random_state=42,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d0bd16e7-1a2b-4ec4-b430-58d7aeaea93b",
"metadata": {},
"outputs": [],
"source": [
"trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
"total_params = sum(p.numel() for p in model.parameters())\n",
"frozen_params = total_params - trainable_params\n",
"\n",
"print(f\"Parameter counts:\")\n",
"print(f\" Trainable: {trainable_params:,} ({100 * trainable_params / total_params:.2f}%)\")\n",
"print(f\" Frozen: {frozen_params:,} ({100 * frozen_params / total_params:.2f}%)\")\n",
"print(f\" Total: {total_params:,}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b14dd996-a869-4c96-a8da-565f414122a8",
"metadata": {},
"outputs": [],
"source": [
"from trl import SFTTrainer\n",
"from transformers import TrainingArguments\n",
"\n",
"training_args = TrainingArguments(\n",
" output_dir=\"./training_output\",\n",
" num_train_epochs=3,\n",
" per_device_train_batch_size=4,\n",
" gradient_accumulation_steps=2,\n",
" learning_rate=2e-4,\n",
" warmup_ratio=0.1,\n",
" logging_steps=10,\n",
" save_strategy=\"no\",\n",
" bf16=True,\n",
" fp16=False,\n",
" optim=\"adamw_8bit\",\n",
" seed=42,\n",
")\n",
"\n",
"print(\"Training configuration:\")\n",
"print(f\" Epochs: {training_args.num_train_epochs}\")\n",
"print(f\" Batch size per device: {training_args.per_device_train_batch_size}\")\n",
"print(f\" Gradient accumulation: {training_args.gradient_accumulation_steps}\")\n",
"print(f\" Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}\")\n",
"print(f\" Learning rate: {training_args.learning_rate}\")\n",
"print(f\" Warmup ratio: {training_args.warmup_ratio}\")\n",
"print(f\" Optimizer: {training_args.optim}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e8783289-1c35-4559-af76-52ad1368ddc5",
"metadata": {},
"outputs": [],
"source": [
"from trl import SFTTrainer\n",
"from transformers import TrainingArguments\n",
"\n",
"training_args = TrainingArguments(\n",
" output_dir=\"./training_output\",\n",
" num_train_epochs=3,\n",
" per_device_train_batch_size=4,\n",
" gradient_accumulation_steps=2,\n",
" learning_rate=2e-4,\n",
" warmup_ratio=0.1,\n",
" logging_steps=10,\n",
" save_strategy=\"no\",\n",
" bf16=True,\n",
" fp16=False,\n",
" optim=\"adamw_8bit\",\n",
" seed=42,\n",
")\n",
"\n",
"print(\"Training configuration:\")\n",
"print(f\" Epochs: {training_args.num_train_epochs}\")\n",
"print(f\" Batch size per device: {training_args.per_device_train_batch_size}\")\n",
"print(f\" Gradient accumulation: {training_args.gradient_accumulation_steps}\")\n",
"print(f\" Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}\")\n",
"print(f\" Learning rate: {training_args.learning_rate}\")\n",
"print(f\" Warmup ratio: {training_args.warmup_ratio}\")\n",
"print(f\" Optimizer: {training_args.optim}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6eb4f82f-1385-450d-8641-a8e57573756d",
"metadata": {},
"outputs": [],
"source": [
"trainer = SFTTrainer(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" train_dataset=dataset,\n",
" args=training_args,\n",
" max_seq_length=512,\n",
" dataset_text_field=\"text\",\n",
" packing=False,\n",
")\n",
"\n",
"print(f\"Trainer created successfully\")\n",
"print(f\" Training examples: {len(dataset)}\")\n",
"print(f\" Steps per epoch: {len(dataset) // (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps)}\")\n",
"print(f\" Total training steps: {trainer.args.max_steps if trainer.args.max_steps > 0 else 'auto'}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fd2849c7-4aea-4cc2-8f0b-46a82ee5fc31",
"metadata": {},
"outputs": [],
"source": [
"print(\"=\" * 60)\n",
"print(\"STARTING TRAINING\")\n",
"print(\"=\" * 60)\n",
"\n",
"trainer_stats = trainer.train()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0a891da0-99ea-4968-8bc2-5c2a511e2d6b",
"metadata": {},
"outputs": [],
"source": [
"print(\"\\nTraining complete!\")\n",
"print(f\" Total steps: {trainer_stats.global_step}\")\n",
"print(f\" Training time: {trainer_stats.metrics['train_runtime']:.1f} seconds\")\n",
"print(f\" Samples per second: {trainer_stats.metrics['train_samples_per_second']:.2f}\")\n",
"print(f\" Final loss: {trainer_stats.metrics['train_loss']:.4f}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63ad7f9a-0fb0-4074-ae8f-eb5490b6ff25",
"metadata": {},
"outputs": [],
"source": [
"OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
"\n",
"model.save_pretrained(OUTPUT_DIR)\n",
"tokenizer.save_pretrained(OUTPUT_DIR)\n",
"\n",
"print(f\"\\nModel saved to: {OUTPUT_DIR}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c3183364-0e6a-4d0f-a43a-83cec929c9e4",
"metadata": {},
"outputs": [],
"source": [
"print(\"Saved files:\")\n",
"total_size = 0\n",
"for f in sorted(OUTPUT_DIR.iterdir()):\n",
" size_kb = f.stat().st_size / 1024\n",
" total_size += size_kb\n",
" print(f\" {f.name}: {size_kb:.1f} KB\")\n",
"print(f\" Total: {total_size / 1024:.1f} MB\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"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.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}