{ "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 }