{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# EasyTranslate: Transformer-Based English-Chinese Translation System\n", "\n", "## Production Entry Point — Person C Integration\n", "\n", "This notebook serves as the primary entry point for the EasyTranslate project.\n", "It handles:\n", "- Environment detection (local vs Google Colab)\n", "- Repository cloning and dependency installation\n", "- Data loading and preprocessing\n", "- Model construction and training\n", "- Evaluation and cloud storage synchronization\n", "\n", "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Environment Detection & Setup\n", "\n", "Detect whether we are running in Google Colab or locally, and configure accordingly." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import sys\n", "import subprocess\n", "import importlib\n", "import json\n", "import shutil\n", "from pathlib import Path\n", "\n", "IN_COLAB = False\n", "try:\n", " import google.colab\n", " IN_COLAB = True\n", "except ImportError:\n", " pass\n", "\n", "print(f\"Running in Google Colab: {IN_COLAB}\")\n", "print(f\"Python version: {sys.version}\")\n", "print(f\"Working directory: {os.getcwd()}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Repository Setup\n", "\n", "Clone the latest code from the repository. In Colab, this pulls fresh code each runtime.\n", "Locally, it ensures the working directory is correct." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── 仓库地址配置 ─────────────────────────────────────────────────────────────\n", "# import os; os.environ[\"EASYTRANSLATE_REPO_URL\"] = \"https://github.com/your-org/your-repo.git\"\n", "REPO_URL = os.environ.get(\n", " \"EASYTRANSLATE_REPO_URL\",\n", " \"https://huggingface.co/sdfjliom/UCAS-EasyTranslate\",\n", ")\n", "_DEFAULT_COLAB_DIR = Path(\"/content/UCAS-EasyTranslate\")\n", "\n", "\n", "def _is_repo_root(path: Path) -> bool:\n", " return (path / \"src\" / \"easytranslate\").exists() and (path / \"setup.py\").exists()\n", "\n", "\n", "def _find_repo_root(start_path: Path):\n", " p = start_path.resolve()\n", " for candidate in [p] + list(p.parents):\n", " if _is_repo_root(candidate):\n", " return candidate\n", " return None\n", "\n", "\n", "resolved_repo = None\n", "\n", "if IN_COLAB:\n", " if _is_repo_root(_DEFAULT_COLAB_DIR):\n", " resolved_repo = _DEFAULT_COLAB_DIR\n", " else:\n", " print(f\"Cloning repository from: {REPO_URL}\")\n", " try:\n", " subprocess.run(\n", " [\"git\", \"lfs\", \"install\"], check=False,\n", " stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n", " )\n", " subprocess.run(\n", " [\"git\", \"clone\", \"--depth\", \"1\", REPO_URL, str(_DEFAULT_COLAB_DIR)],\n", " check=True,\n", " stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,\n", " )\n", " resolved_repo = _DEFAULT_COLAB_DIR\n", " print(f\"Cloned to: {_DEFAULT_COLAB_DIR}\")\n", " except subprocess.CalledProcessError as e:\n", " print(f\"Clone failed:\\n{e.stdout}\")\n", "\n", " # Fallback: search Drive or current dir\n", " if resolved_repo is None:\n", " for candidate in [\n", " Path.cwd(),\n", " Path(\"/content\"),\n", " Path(\"/content/drive/MyDrive/UCAS-EasyTranslate\"),\n", " Path(\"/content/drive/MyDrive/Colab Notebooks/UCAS-EasyTranslate\"),\n", " ]:\n", " root = _find_repo_root(candidate)\n", " if root:\n", " resolved_repo = root\n", " print(f\"Found existing repo at: {resolved_repo}\")\n", " break\n", "\n", " if resolved_repo is None:\n", " raise FileNotFoundError(\n", " \"Cannot locate UCAS-EasyTranslate repository.\\n\"\n", " \"Options:\\n\"\n", " \" (1) Set EASYTRANSLATE_REPO_URL to a publicly accessible git URL.\\n\"\n", " \" (2) Manually clone to /content/UCAS-EasyTranslate.\\n\"\n", " \" (3) Place repo in Google Drive and mount Drive first.\"\n", " )\n", "else:\n", " resolved_repo = _find_repo_root(Path.cwd()) or Path.cwd()\n", " print(f\"Using local repository at: {resolved_repo}\")\n", "\n", "# Always keep REPO_DIR in sync with wherever the repo actually is\n", "REPO_DIR = Path(resolved_repo)\n", "os.chdir(REPO_DIR)\n", "sys.path.insert(0, str(REPO_DIR / \"src\"))\n", "print(f\"Repository directory: {REPO_DIR}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Dependency Installation\n", "\n", "Install all required packages. In Colab, PyTorch is pre-installed." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "if IN_COLAB:\n", " # 1. Only install packages that Colab does NOT ship.\n", " # Do NOT touch torch/numpy/pandas — Colab's preinstalled versions are fine.\n", " %pip install -q --upgrade pip setuptools wheel\n", "\n", " # Core HuggingFace packages (Colab may have older versions)\n", " %pip install -q \"transformers>=4.36.0,<4.45.0\" \\\n", " \"datasets>=2.16.0,<3.0.0\" \\\n", " \"tokenizers>=0.15.0,<0.20.0\" \\\n", " \"sentencepiece>=0.2.0\" \\\n", " \"accelerate>=0.25.0,<0.35.0\" \\\n", " \"peft>=0.7.0,<0.12.0\"\n", "\n", " # Evaluation / config packages\n", " %pip install -q \"sacrebleu>=2.4.0\" \\\n", " \"omegaconf>=2.3.0,<3.0.0\" \\\n", " \"rich>=13.0.0\"\n", "\n", " # 2. Install project code without re-resolving heavy dependencies.\n", " # (torch, numpy, etc. are already present from Colab runtime.)\n", " %pip install -q --no-deps -e .\n", "\n", " print(\"All packages installed successfully.\")\n", "else:\n", " %pip install -q --upgrade pip setuptools wheel\n", " %pip install -q -r requirements.txt\n", " %pip install -q -e .\n", " print(\"Dependencies installed successfully.\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. GPU Verification\n", "\n", "Verify GPU availability and display hardware information." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "print(f\"PyTorch version: {torch.__version__}\")\n", "print(f\"CUDA available: {torch.cuda.is_available()}\")\n", "\n", "if torch.cuda.is_available():\n", " print(f\"CUDA version: {torch.version.cuda}\")\n", " print(f\"GPU count: {torch.cuda.device_count()}\")\n", " for i in range(torch.cuda.device_count()):\n", " print(f\" GPU {i}: {torch.cuda.get_device_name(i)}\")\n", " props = torch.cuda.get_device_properties(i)\n", " print(f\" Memory: {props.total_memory / 1024**3:.1f} GB\")\n", " print(f\" Compute Capability: {props.major}.{props.minor}\")\n", "else:\n", " print(\"WARNING: No GPU detected. Training will be very slow on CPU.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Google Drive Mount (Colab Only)\n", "\n", "Mount Google Drive for persistent storage of checkpoints and results." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "DRIVE_MOUNTED = False\n", "DRIVE_BASE = \"/content/drive/MyDrive/EasyTranslate\"\n", "\n", "if IN_COLAB:\n", " from google.colab import drive\n", " drive.mount(\"/content/drive\")\n", " DRIVE_MOUNTED = os.path.exists(\"/content/drive\")\n", " if DRIVE_MOUNTED:\n", " os.makedirs(DRIVE_BASE, exist_ok=True)\n", " print(f\"Google Drive mounted. Base path: {DRIVE_BASE}\")\n", " else:\n", " print(\"WARNING: Google Drive mount failed\")\n", "else:\n", " print(\"Not in Colab, skipping Google Drive mount\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6. Configuration Loading\n", "\n", "Load and display the project configuration." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "\n", "print(f\"NumPy version : {np.__version__}\")\n", "print(f\"PyTorch version: {torch.__version__}\")\n", "\n", "from easytranslate.utils.config import load_config, config_to_dict\n", "from easytranslate.utils.seed import set_seed\n", "from easytranslate.utils.logging import setup_logging\n", "\n", "config = load_config(\"configs/default_config.yaml\")\n", "config_dict = config_to_dict(config)\n", "\n", "exp_cfg = config_dict.get(\"experiment\", {})\n", "seed = exp_cfg.get(\"seed\", 42)\n", "set_seed(seed)\n", "\n", "log_cfg = config_dict.get(\"logging\", {})\n", "setup_logging(\n", " log_dir=log_cfg.get(\"log_dir\", \"logs/\"),\n", " log_file=\"easytranslate.log\",\n", ")\n", "\n", "# ── Colab runtime overrides ───────────────────────────────────────────────────\n", "if IN_COLAB:\n", " train_cfg = config_dict.setdefault(\"training\", {})\n", "\n", " # Mixed precision: use bf16 on Ampere+ GPUs, fp16 otherwise, none on CPU\n", " if torch.cuda.is_available():\n", " gpu_cap = torch.cuda.get_device_capability(0)\n", " if gpu_cap[0] >= 8: # A100/A10 → bf16\n", " train_cfg[\"fp16\"] = False\n", " train_cfg[\"bf16\"] = True\n", " else: # T4/P100/V100 → fp16\n", " train_cfg[\"fp16\"] = True\n", " train_cfg[\"bf16\"] = False\n", " else:\n", " train_cfg[\"fp16\"] = False\n", " train_cfg[\"bf16\"] = False\n", "\n", " # Reduce epochs for a Colab demo run\n", " train_cfg.setdefault(\"epochs\", 10)\n", "\n", " # Colab-friendly batch / gradient accumulation settings\n", " train_cfg.setdefault(\"batch_size\", 32)\n", " train_cfg.setdefault(\"gradient_accumulation_steps\", 4)\n", "\n", "print(f\"Configuration loaded. Experiment seed: {seed}\")\n", "print(f\"Model type : {config_dict['model']['type']}\")\n", "print(f\"Training : epochs={config_dict['training']['epochs']}, \"\n", " f\"fp16={config_dict['training']['fp16']}, \"\n", " f\"bf16={config_dict['training']['bf16']}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 7. Data Loading & Preprocessing\n", "\n", "Load the WMT19 zh-en dataset, train the BPE tokenizer, and prepare DataLoaders." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from easytranslate.data import (\n", " TranslationDataset,\n", " TranslationCollator,\n", " DynamicBatchSampler,\n", " build_tokenizer,\n", " load_wmt_dataset,\n", " preprocess_pipeline,\n", ")\n", "from torch.utils.data import DataLoader\n", "\n", "data_cfg = config_dict.get(\"data\", {})\n", "preproc_cfg = data_cfg.get(\"preprocessing\", {})\n", "loader_cfg = data_cfg.get(\"dataloader\", {})\n", "\n", "# ── Cap dataset sizes for Colab to avoid RAM/time issues ─────────────────────\n", "MAX_TRAIN_SAMPLES = 200_000 if IN_COLAB else None # None → use full dataset\n", "MAX_VAL_SAMPLES = 5_000 if IN_COLAB else None\n", "\n", "print(\"Loading WMT19 zh-en dataset (this may take several minutes on first run)...\")\n", "raw_dataset = load_wmt_dataset(\n", " year=data_cfg.get(\"wmt\", {}).get(\"year\", \"19\"),\n", " language_pair=data_cfg.get(\"wmt\", {}).get(\"language_pair\", \"zh-en\"),\n", " src_lang=\"en\",\n", " tgt_lang=\"zh\",\n", ")\n", "\n", "train_raw = raw_dataset[\"train\"]\n", "\n", "# Prefer \"validation\", fall back to \"dev\", then use a small slice of train\n", "_val_split_name = next(\n", " (k for k in (\"validation\", \"dev\", \"valid\") if k in raw_dataset),\n", " None,\n", ")\n", "val_raw = raw_dataset[_val_split_name] if _val_split_name else train_raw\n", "\n", "print(f\"Raw training samples : {len(train_raw['src'])}\")\n", "print(f\"Raw validation samples: {len(val_raw['src'])} \"\n", " f\"(split='{_val_split_name or 'train (fallback)'}')\")\n", "\n", "# Apply sample caps BEFORE preprocessing to save time\n", "train_src_raw = train_raw[\"src\"][:MAX_TRAIN_SAMPLES] if MAX_TRAIN_SAMPLES else list(train_raw[\"src\"])\n", "train_tgt_raw = train_raw[\"tgt\"][:MAX_TRAIN_SAMPLES] if MAX_TRAIN_SAMPLES else list(train_raw[\"tgt\"])\n", "val_src_raw = val_raw[\"src\"][:MAX_VAL_SAMPLES] if MAX_VAL_SAMPLES else list(val_raw[\"src\"])\n", "val_tgt_raw = val_raw[\"tgt\"][:MAX_VAL_SAMPLES] if MAX_VAL_SAMPLES else list(val_raw[\"tgt\"])\n", "\n", "print(f\"\\nPreprocessing training data ({len(train_src_raw)} samples)...\")\n", "train_src, train_tgt = preprocess_pipeline(\n", " train_src_raw, train_tgt_raw,\n", " max_src_len=preproc_cfg.get(\"max_src_len\", 256),\n", " max_tgt_len=preproc_cfg.get(\"max_tgt_len\", 256),\n", " filter_by_length_enabled=preproc_cfg.get(\"filter_by_length\", True),\n", " length_ratio_threshold=preproc_cfg.get(\"length_ratio_threshold\", 3.0),\n", ")\n", "print(f\"Preprocessed training samples : {len(train_src)}\")\n", "\n", "print(f\"Preprocessing validation data ({len(val_src_raw)} samples)...\")\n", "val_src, val_tgt = preprocess_pipeline(\n", " val_src_raw, val_tgt_raw,\n", " max_src_len=preproc_cfg.get(\"max_src_len\", 256),\n", " max_tgt_len=preproc_cfg.get(\"max_tgt_len\", 256),\n", " filter_by_length_enabled=preproc_cfg.get(\"filter_by_length\", True),\n", " length_ratio_threshold=preproc_cfg.get(\"length_ratio_threshold\", 3.0),\n", ")\n", "print(f\"Preprocessed validation samples: {len(val_src)}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 8. Tokenizer Training\n", "\n", "Train a byte-level BPE tokenizer on the combined source and target texts." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "tok_cfg = config_dict.get(\"tokenizer\", {})\n", "\n", "print(\"Building tokenizer...\")\n", "all_train_texts = train_src + train_tgt\n", "tokenizer = build_tokenizer(tok_cfg, train_texts=all_train_texts)\n", "\n", "print(f\"Tokenizer vocabulary size: {tokenizer.vocab_size}\")\n", "print(f\"Special tokens: PAD={tokenizer.pad_token_id}, BOS={tokenizer.bos_token_id}, EOS={tokenizer.eos_token_id}\")\n", "\n", "test_encode = tokenizer.encode(\"Hello world\", add_special_tokens=True)\n", "test_decode = tokenizer.decode(test_encode)\n", "print(f\"Encode test: {test_encode[:10]}...\")\n", "print(f\"Decode test: {test_decode[:50]}...\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 9. Dataset & DataLoader Construction\n", "\n", "Build PyTorch datasets and dataloaders with dynamic batching." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "max_src_len = preproc_cfg.get(\"max_src_len\", 256)\n", "max_tgt_len = preproc_cfg.get(\"max_tgt_len\", 256)\n", "\n", "train_dataset = TranslationDataset(\n", " train_src, train_tgt,\n", " tokenizer=tokenizer,\n", " max_src_len=max_src_len,\n", " max_tgt_len=max_tgt_len,\n", ")\n", "val_dataset = TranslationDataset(\n", " val_src, val_tgt,\n", " tokenizer=tokenizer,\n", " max_src_len=max_src_len,\n", " max_tgt_len=max_tgt_len,\n", ")\n", "\n", "collator = TranslationCollator(\n", " pad_token_id=tokenizer.pad_token_id,\n", " label_pad_token_id=-100,\n", ")\n", "\n", "loader_cfg = config_dict.get(\"data\", {}).get(\"dataloader\", {})\n", "batch_size = loader_cfg.get(\"batch_size\", 32)\n", "# 2 workers on GPU Colab; 0 on CPU (no multiprocessing overhead)\n", "num_workers = 2 if IN_COLAB and torch.cuda.is_available() else 0\n", "# Disable dynamic batching on Colab: computing exact token lengths for 200k\n", "# sentences requires a full tokenizer pass and can take 30+ minutes.\n", "use_dynamic = loader_cfg.get(\"dynamic_batching\", True) and not IN_COLAB\n", "\n", "if use_dynamic:\n", " max_tokens = loader_cfg.get(\"max_tokens_per_batch\", 8192)\n", " print(\"Computing sequence lengths for dynamic batching (approximate)...\")\n", "\n", " def _approx_len(src_text: str, tgt_text: str) -> int:\n", " \"\"\"Fast character-based length estimate — no tokenizer call needed.\"\"\"\n", " def _tok_est(t: str) -> int:\n", " cjk = sum(1 for c in t if \"\\u4e00\" <= c <= \"\\u9fff\")\n", " return (len(t) - cjk) // 4 + cjk + 2 # rough BPE estimate\n", " return max(_tok_est(src_text), _tok_est(tgt_text))\n", "\n", " train_lengths = [_approx_len(s, t) for s, t in zip(train_src, train_tgt)]\n", " train_sampler = DynamicBatchSampler(\n", " train_lengths,\n", " max_tokens_per_batch=max_tokens,\n", " shuffle=True,\n", " )\n", " train_loader = DataLoader(\n", " train_dataset,\n", " batch_sampler=train_sampler,\n", " collate_fn=collator,\n", " num_workers=num_workers,\n", " pin_memory=torch.cuda.is_available(),\n", " )\n", "else:\n", " if IN_COLAB and loader_cfg.get(\"dynamic_batching\", True):\n", " print(\"Note: dynamic batching disabled on Colab (would require full tokenizer pass on all samples).\")\n", " train_loader = DataLoader(\n", " train_dataset,\n", " batch_size=batch_size,\n", " shuffle=True,\n", " collate_fn=collator,\n", " num_workers=num_workers,\n", " pin_memory=torch.cuda.is_available(),\n", " )\n", "\n", "val_loader = DataLoader(\n", " val_dataset,\n", " batch_size=batch_size,\n", " shuffle=False,\n", " collate_fn=collator,\n", " num_workers=num_workers,\n", " pin_memory=torch.cuda.is_available(),\n", ")\n", "\n", "print(f\"Training batches : ~{len(train_loader)}\")\n", "print(f\"Validation batches: {len(val_loader)}\")\n", "\n", "sample_batch = next(iter(train_loader))\n", "print(\"Sample batch shapes:\")\n", "for k, v in sample_batch.items():\n", " if isinstance(v, torch.Tensor):\n", " print(f\" {k}: {list(v.shape)}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 10. Model Construction\n", "\n", "Build the Transformer model based on configuration." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from easytranslate.model import TransformerTranslationModel\n", "\n", "model_cfg = config_dict.get(\"model\", {})\n", "model_type = model_cfg.get(\"type\", \"transformer_scratch\")\n", "\n", "if model_type == \"transformer_scratch\":\n", " tf_cfg = dict(model_cfg.get(\"transformer\", {}))\n", "\n", " # ── Colab quick-run: smaller model to fit in Colab RAM/VRAM ──────────────\n", " # Default full model: d_model=512, 6 enc/dec layers (~75 M params)\n", " # Colab quick model: d_model=256, 3 enc/dec layers (~12 M params)\n", " # Set COLAB_FULL_MODEL=1 in env to skip this override.\n", " if IN_COLAB and not os.environ.get(\"COLAB_FULL_MODEL\"):\n", " tf_cfg.setdefault(\"d_model\", 256)\n", " tf_cfg.setdefault(\"nhead\", 4)\n", " tf_cfg.setdefault(\"num_encoder_layers\", 3)\n", " tf_cfg.setdefault(\"num_decoder_layers\", 3)\n", " tf_cfg.setdefault(\"dim_feedforward\", 1024)\n", " print(\"Colab mode: using compact model (d_model=256, 3 layers).\")\n", " print(\"To use the full model, run: import os; os.environ['COLAB_FULL_MODEL']='1'\")\n", "\n", " model = TransformerTranslationModel(\n", " src_vocab_size=tokenizer.vocab_size,\n", " tgt_vocab_size=tokenizer.vocab_size,\n", " d_model=tf_cfg.get(\"d_model\", 512),\n", " nhead=tf_cfg.get(\"nhead\", 8),\n", " num_encoder_layers=tf_cfg.get(\"num_encoder_layers\", 6),\n", " num_decoder_layers=tf_cfg.get(\"num_decoder_layers\", 6),\n", " dim_feedforward=tf_cfg.get(\"dim_feedforward\", 2048),\n", " dropout=tf_cfg.get(\"dropout\", 0.1),\n", " activation=tf_cfg.get(\"activation\", \"gelu\"),\n", " max_seq_len=tf_cfg.get(\"max_seq_len\", 512),\n", " use_flash_attention=tf_cfg.get(\"use_flash_attention\", True),\n", " use_rotary_embedding=tf_cfg.get(\"use_rotary_embedding\", True),\n", " pre_norm=tf_cfg.get(\"pre_norm\", True),\n", " pad_id=tokenizer.pad_token_id,\n", " share_embedding=False,\n", " )\n", " print(\"Built Transformer from scratch\")\n", "\n", "elif model_type in (\"finetune_nllb\", \"finetune_mbart\"):\n", " from easytranslate.model.finetune import load_pretrained_model, setup_lora\n", " pt_cfg = model_cfg.get(\"pretrained\", {})\n", " model, hf_tokenizer = load_pretrained_model(\n", " model_name=pt_cfg.get(\"model_name\", \"facebook/nllb-200-distilled-600M\"),\n", " src_lang=pt_cfg.get(\"src_lang\", \"eng_Latn\"),\n", " tgt_lang=pt_cfg.get(\"tgt_lang\", \"zho_Hans\"),\n", " )\n", " if pt_cfg.get(\"use_lora\", True):\n", " lora_cfg = pt_cfg.get(\"lora\", {})\n", " model = setup_lora(\n", " model,\n", " r=lora_cfg.get(\"r\", 16),\n", " alpha=lora_cfg.get(\"alpha\", 32),\n", " dropout=lora_cfg.get(\"dropout\", 0.05),\n", " target_modules=lora_cfg.get(\"target_modules\", [\"q_proj\", \"v_proj\"]),\n", " )\n", " print(f\"Loaded pretrained model: {pt_cfg.get('model_name')}\")\n", "\n", "else:\n", " raise ValueError(f\"Unknown model type: {model_type}\")\n", "\n", "total_params = sum(p.numel() for p in model.parameters())\n", "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n", "print(f\"Total parameters : {total_params:,}\")\n", "print(f\"Trainable parameters : {trainable_params:,}\")\n", "print(f\"Trainable ratio : {100 * trainable_params / total_params:.2f}%\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11. Quick Forward Pass Test\n", "\n", "Verify the model can perform a forward pass with correct output dimensions." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "model = model.to(device)\n", "model.eval()\n", "\n", "test_batch = next(iter(train_loader))\n", "test_src = test_batch[\"src_ids\"][:2].to(device)\n", "test_tgt = test_batch[\"tgt_input_ids\"][:2].to(device)\n", "test_src_mask = test_batch[\"src_padding_mask\"][:2].to(device)\n", "test_tgt_mask = test_batch[\"tgt_padding_mask\"][:2].to(device)\n", "\n", "with torch.no_grad():\n", " logits = model(test_src, test_tgt, test_src_mask, test_tgt_mask)\n", "\n", "print(f\"Input src shape: {list(test_src.shape)}\")\n", "print(f\"Input tgt shape: {list(test_tgt.shape)}\")\n", "print(f\"Output logits shape: {list(logits.shape)}\")\n", "print(f\"Expected output shape: [B, T, vocab_size] = [{test_tgt.size(0)}, {test_tgt.size(1)}, {tokenizer.vocab_size}]\")\n", "assert logits.size(-1) == tokenizer.vocab_size, f\"Vocab size mismatch: {logits.size(-1)} vs {tokenizer.vocab_size}\"\n", "print(\"Forward pass test PASSED\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 12. Training Execution\n", "\n", "This cell DEFINES the training setup but does NOT execute training automatically.\n", "To start training, run the cell below this one." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# sacrebleu is now installed in the dependency cell above.\n", "# This cell just verifies it is importable before we initialize the Trainer.\n", "try:\n", " importlib.import_module(\"sacrebleu\")\n", " print(\"sacrebleu OK\")\n", "except ImportError:\n", " print(\"sacrebleu missing — installing now...\")\n", " subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"sacrebleu>=2.4.0\"], check=True)\n", "\n", "from easytranslate.training import Trainer\n", "from easytranslate.evaluation import Evaluator\n", "\n", "evaluator = Evaluator(\n", " model=model,\n", " tokenizer=tokenizer,\n", " config=config_dict,\n", ")\n", "\n", "trainer = Trainer(\n", " model=model,\n", " train_loader=train_loader,\n", " val_loader=val_loader,\n", " config=config_dict,\n", " evaluator=evaluator,\n", ")\n", "\n", "# Output / plot directories live inside the repo\n", "OUTPUT_DIR = REPO_DIR / \"outputs\"\n", "PLOTS_DIR = OUTPUT_DIR / \"plots\"\n", "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", "PLOTS_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "print(\"Trainer initialized successfully\")\n", "print(f\" Device : {trainer.device}\")\n", "print(f\" FP16 / BF16 : {trainer.fp16} / {trainer.bf16}\")\n", "print(f\" Gradient accumulation steps: {trainer.gradient_accumulation_steps}\")\n", "print(f\" Number of epochs : {trainer.num_epochs}\")\n", "print(f\" Checkpoint directory : {trainer.checkpoint_dir}\")\n", "print(f\" Output directory : {OUTPUT_DIR}\")\n", "print()\n", "print(\">>> Run the NEXT cell to start training.\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 13. Start Training\n", "\n", "**Run this cell to begin training.** This will execute the full training loop.\n", "Training progress will be displayed via tqdm progress bars." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "START_TRAINING = True\n", "\n", "if START_TRAINING:\n", " print(\"=\" * 60)\n", " print(\" Starting Training...\")\n", " print(\"=\" * 60)\n", " trainer.train()\n", "else:\n", " print(\"Training skipped. Set START_TRAINING = True to begin.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 14. Evaluation on Test Set\n", "\n", "After training completes, evaluate the best model on the validation set." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "EVAL_RESULTS = {}\n", "\n", "best_ckpt = trainer.checkpoint_dir / \"best_model.pt\"\n", "\n", "if best_ckpt.exists():\n", " print(f\"Loading best model from {best_ckpt}\")\n", " checkpoint = torch.load(best_ckpt, map_location=device, weights_only=True)\n", " model.load_state_dict(checkpoint[\"model_state_dict\"])\n", " model = model.to(device)\n", " model.eval()\n", "\n", " evaluator = Evaluator(model=model, tokenizer=tokenizer, config=config_dict)\n", "\n", " print(\"Running evaluation on validation set...\")\n", " EVAL_RESULTS = evaluator.evaluate(\n", " val_loader,\n", " src_texts=val_src,\n", " ref_texts=val_tgt,\n", " )\n", "\n", " print(\"\\n\" + \"=\" * 60)\n", " print(\" Evaluation Results\")\n", " print(\"=\" * 60)\n", " for metric, score in EVAL_RESULTS.items():\n", " if isinstance(score, (int, float)):\n", " print(f\" {metric:>12s}: {score:.4f}\")\n", "\n", " eval_path = OUTPUT_DIR / \"evaluation_results.json\"\n", " with open(eval_path, \"w\", encoding=\"utf-8\") as f:\n", " json.dump(EVAL_RESULTS, f, indent=2, ensure_ascii=False)\n", " print(f\"\\nSaved evaluation results → {eval_path}\")\n", "else:\n", " print(\"No best_model.pt found. Run training first (Cell 13).\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 15. Translation Demo\n", "\n", "Test the trained model with some example translations." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "test_sentences = [\n", " \"Hello, how are you today?\",\n", " \"Machine translation is an important field of natural language processing.\",\n", " \"The weather is beautiful and I want to go for a walk.\",\n", " \"Deep learning has revolutionized artificial intelligence research.\",\n", "]\n", "\n", "TRANSLATION_RESULTS = []\n", "\n", "if best_ckpt.exists():\n", " print(\"Translating example sentences...\")\n", " print(\"-\" * 60)\n", " for sentence in test_sentences:\n", " translation = evaluator.translate_single(sentence)\n", " TRANSLATION_RESULTS.append({\"source_en\": sentence, \"target_zh\": translation})\n", " print(f\"[EN] {sentence}\")\n", " print(f\"[ZH] {translation}\")\n", " print()\n", "\n", " translation_path = OUTPUT_DIR / \"translation_examples.json\"\n", " with open(translation_path, \"w\", encoding=\"utf-8\") as f:\n", " json.dump(TRANSLATION_RESULTS, f, indent=2, ensure_ascii=False)\n", " print(f\"Saved translation examples → {translation_path}\")\n", "else:\n", " print(\"No trained model checkpoint available. Run training first (Cell 13).\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 16. Cloud Storage Synchronization\n", "\n", "Sync all training artifacts (checkpoints, logs, summaries) to Google Drive." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from easytranslate.utils.cloud_storage import sync_all_to_drive\n", "\n", "# Use trainer.log_dir if available, otherwise fall back to a sensible default\n", "_log_dir = getattr(trainer, \"log_dir\", REPO_DIR / \"logs\")\n", "\n", "if IN_COLAB and DRIVE_MOUNTED:\n", " print(\"Syncing training artifacts to Google Drive...\")\n", " sync_results = sync_all_to_drive(\n", " checkpoint_dir=str(trainer.checkpoint_dir),\n", " log_dir=str(_log_dir),\n", " drive_base_path=DRIVE_BASE,\n", " )\n", "\n", " drive_outputs_dir = Path(DRIVE_BASE) / \"outputs\"\n", " drive_outputs_dir.mkdir(parents=True, exist_ok=True)\n", "\n", " for artifact_file in [\n", " OUTPUT_DIR / \"evaluation_results.json\",\n", " OUTPUT_DIR / \"translation_examples.json\",\n", " OUTPUT_DIR / \"training_report.json\",\n", " ]:\n", " if artifact_file.exists():\n", " shutil.copy2(artifact_file, drive_outputs_dir / artifact_file.name)\n", " print(f\" Copied {artifact_file.name} → Drive\")\n", "\n", " if PLOTS_DIR.exists():\n", " drive_plots_dir = drive_outputs_dir / \"plots\"\n", " drive_plots_dir.mkdir(parents=True, exist_ok=True)\n", " for png_file in PLOTS_DIR.glob(\"*.png\"):\n", " shutil.copy2(png_file, drive_plots_dir / png_file.name)\n", " print(f\" Copied plot {png_file.name} → Drive\")\n", "\n", " print(f\"\\nSync complete. Drive base: {DRIVE_BASE}\")\n", " print(f\"Sync details: {sync_results}\")\n", "\n", "elif not IN_COLAB:\n", " print(\"Running locally. Artifacts are already on disk:\")\n", " print(f\" Checkpoints : {trainer.checkpoint_dir}\")\n", " print(f\" Logs : {_log_dir}\")\n", " print(f\" Outputs : {OUTPUT_DIR}\")\n", "else:\n", " print(\"Google Drive not mounted. Artifacts saved locally only.\")\n", " print(\"Mount Drive (Cell 5) and re-run this cell to sync results.\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 17. Training Summary\n", "\n", "Display the final training summary including loss curves and best metrics." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "TRAINING_SUMMARY = {}\n", "summary_path = trainer.checkpoint_dir / \"training_summary.json\"\n", "\n", "if summary_path.exists():\n", " with open(summary_path, \"r\", encoding=\"utf-8\") as f:\n", " TRAINING_SUMMARY = json.load(f)\n", "\n", " print(\"=\" * 60)\n", " print(\" Training Summary\")\n", " print(\"=\" * 60)\n", " print(f\" Best epoch : {TRAINING_SUMMARY.get('best_epoch', 'N/A')}\")\n", " print(f\" Best metric : {TRAINING_SUMMARY.get('metric_name', 'N/A')} = \"\n", " f\"{TRAINING_SUMMARY.get('best_metric', 'N/A')}\")\n", " print(f\" Total steps : {TRAINING_SUMMARY.get('total_steps', 'N/A')}\")\n", "\n", " losses = TRAINING_SUMMARY.get(\"train_loss_history\", [])\n", " if losses:\n", " print(f\" Initial loss: {losses[0]:.4f}\")\n", " print(f\" Final loss : {losses[-1]:.4f}\")\n", " print(f\" Reduction : {losses[0] - losses[-1]:.4f}\")\n", "\n", " # Merge all results into one report file\n", " report = {\n", " \"training_summary\": TRAINING_SUMMARY,\n", " \"evaluation_results\": EVAL_RESULTS,\n", " \"translation_examples\": TRANSLATION_RESULTS,\n", " }\n", " report_path = OUTPUT_DIR / \"training_report.json\"\n", " with open(report_path, \"w\", encoding=\"utf-8\") as f:\n", " json.dump(report, f, indent=2, ensure_ascii=False)\n", " print(f\"\\nMerged report saved → {report_path}\")\n", "else:\n", " print(\"Training summary not yet available. Complete training first.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "PLOTS_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "if not TRAINING_SUMMARY:\n", " print(\"No training summary found. Run training and summary cells first.\")\n", "else:\n", " train_losses = TRAINING_SUMMARY.get(\"train_loss_history\", [])\n", " val_history = TRAINING_SUMMARY.get(\"val_metrics_history\", [])\n", "\n", " # 1) Train Loss Curve\n", " if train_losses:\n", " epochs = list(range(1, len(train_losses) + 1))\n", " plt.figure(figsize=(8, 5))\n", " plt.plot(epochs, train_losses, marker=\"o\", linewidth=2)\n", " plt.title(\"Training Loss by Epoch\")\n", " plt.xlabel(\"Epoch\")\n", " plt.ylabel(\"Loss\")\n", " plt.grid(alpha=0.3)\n", " loss_plot_path = PLOTS_DIR / \"train_loss_curve.png\"\n", " plt.tight_layout()\n", " plt.savefig(loss_plot_path, dpi=180)\n", " plt.show()\n", " print(f\"Saved plot: {loss_plot_path}\")\n", "\n", " # 2) Validation Metrics Curves\n", " if val_history:\n", " metric_keys = sorted({k for m in val_history for k in m.keys() if isinstance(m.get(k), (int, float))})\n", " metric_keys = [k for k in metric_keys if k != \"val_loss\"]\n", "\n", " if metric_keys:\n", " n = len(metric_keys)\n", " rows = (n + 1) // 2\n", " plt.figure(figsize=(12, max(4, rows * 3.5)))\n", " for i, key in enumerate(metric_keys, start=1):\n", " vals = [m.get(key, None) for m in val_history]\n", " xs = [idx + 1 for idx, v in enumerate(vals) if v is not None]\n", " ys = [v for v in vals if v is not None]\n", " if not ys:\n", " continue\n", " plt.subplot(rows, 2, i)\n", " plt.plot(xs, ys, marker=\"o\", linewidth=1.8)\n", " plt.title(key)\n", " plt.xlabel(\"Epoch\")\n", " plt.ylabel(key)\n", " plt.grid(alpha=0.3)\n", "\n", " metrics_plot_path = PLOTS_DIR / \"validation_metrics_curves.png\"\n", " plt.tight_layout()\n", " plt.savefig(metrics_plot_path, dpi=180)\n", " plt.show()\n", " print(f\"Saved plot: {metrics_plot_path}\")\n", "\n", " # 3) Final Evaluation Bar Chart\n", " if \"EVAL_RESULTS\" in globals() and EVAL_RESULTS:\n", " scalar_items = {k: v for k, v in EVAL_RESULTS.items() if isinstance(v, (int, float))}\n", " if scalar_items:\n", " names = list(scalar_items.keys())\n", " values = [scalar_items[k] for k in names]\n", " plt.figure(figsize=(10, 5))\n", " bars = plt.bar(names, values)\n", " plt.title(\"Final Evaluation Metrics\")\n", " plt.ylabel(\"Score\")\n", " plt.xticks(rotation=30)\n", " plt.grid(axis=\"y\", alpha=0.25)\n", " for bar, val in zip(bars, values):\n", " plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height(), f\"{val:.3f}\", ha=\"center\", va=\"bottom\", fontsize=9)\n", " eval_plot_path = PLOTS_DIR / \"final_evaluation_metrics.png\"\n", " plt.tight_layout()\n", " plt.savefig(eval_plot_path, dpi=180)\n", " plt.show()\n", " print(f\"Saved plot: {eval_plot_path}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "## Appendix: Module Architecture Overview\n", "\n", "```\n", "EasyTranslate System Architecture\n", "=================================\n", "\n", "Entry Point: EasyTranslate_Production.ipynb (this notebook)\n", " |\n", " +-- Environment Detection (Colab vs Local)\n", " +-- Repository Cloning (git clone/pull)\n", " +-- Dependency Installation\n", " |\n", " +-- Configuration Layer [utils/config.py]\n", " | +-- load_config() : YAML -> OmegaConf DictConfig\n", " | +-- merge_configs() : CLI overrides merge\n", " | +-- config_from_cli() : Full config pipeline\n", " |\n", " +-- Data Layer [data/] — Person A\n", " | +-- load_wmt_dataset() : HuggingFace datasets loader\n", " | +-- preprocess_pipeline() : Clean + filter + deduplicate\n", " | +-- build_tokenizer() : BPE / pretrained tokenizer\n", " | +-- TranslationDataset() : PyTorch Dataset\n", " | +-- TranslationCollator() : Padding + mask generation\n", " | +-- DynamicBatchSampler() : Token-budget batching\n", " |\n", " +-- Model Layer [model/] — Person B\n", " | +-- TransformerTranslationModel() : Full Enc-Dec model\n", " | +-- TransformerEncoder() : N-layer encoder\n", " | +-- TransformerDecoder() : N-layer decoder\n", " | +-- FlashMultiHeadAttention() : Flash Attention 2\n", " | +-- RotaryPositionalEmbedding() : RoPE encoding\n", " | +-- load_pretrained_model() : NLLB/mBART loader\n", " | +-- setup_lora() : LoRA configuration\n", " |\n", " +-- Training Layer [training/] — Person C\n", " | +-- Trainer() : Full training controller\n", " | | +-- _train_one_epoch() : Mixed precision loop\n", " | | +-- _validate() : Validation loop\n", " | | +-- _save_checkpoint() : Checkpoint persistence\n", " | | +-- _load_checkpoint() : Resume training\n", " | | +-- _should_early_stop() : Early stopping logic\n", " | | +-- _setup_distributed() : DDP/DeepSpeed setup\n", " | +-- LabelSmoothedCrossEntropyLoss() : Label smoothing loss\n", " | +-- build_optimizer() : AdamW with param groups\n", " | +-- build_scheduler() : Cosine/InverseSqrt/LR\n", " |\n", " +-- Evaluation Layer [evaluation/] — Person D\n", " | +-- Evaluator() : Unified evaluation interface\n", " | +-- greedy_decode() : Greedy decoding\n", " | +-- beam_search_decode() : Beam search decoding\n", " | +-- sample_decode() : Sampling (temp/top-k/top-p)\n", " | +-- compute_bleu() : SacreBLEU metric\n", " | +-- compute_comet() : COMET neural metric\n", " | +-- compute_chrf() : chrF++ metric\n", " |\n", " +-- Cloud Storage [utils/cloud_storage.py] — Person C\n", " +-- is_colab_environment() : Environment detection\n", " +-- mount_google_drive() : Drive authentication\n", " +-- sync_checkpoints_to_drive() : Checkpoint backup\n", " +-- sync_logs_to_drive() : Log backup\n", " +-- sync_all_to_drive() : Full sync pipeline\n", "```\n", "\n", "## Module Interface Contracts\n", "\n", "### Tokenizer Interface (Person A -> B, C, D)\n", "```python\n", "tokenizer.encode(text: str) -> list[int]\n", "tokenizer.decode(ids: list[int]) -> str\n", "tokenizer.vocab_size -> int\n", "tokenizer.pad_token_id -> int\n", "tokenizer.bos_token_id -> int\n", "tokenizer.eos_token_id -> int\n", "```\n", "\n", "### Model Interface (Person B -> C, D)\n", "```python\n", "# Training forward pass\n", "logits = model(src_ids, tgt_input_ids, src_padding_mask, tgt_padding_mask)\n", "# logits: [B, T, vocab_size]\n", "\n", "# Inference\n", "encoder_output = model.encode(src_ids, src_padding_mask)\n", "next_logits = model.decode_step(tgt_input_ids, encoder_output, src_padding_mask)\n", "```\n", "\n", "### Batch Format (Person A -> C)\n", "```python\n", "batch = {\n", " \"src_ids\": Tensor[B, S],\n", " \"tgt_input_ids\": Tensor[B, T],\n", " \"labels\": Tensor[B, T],\n", " \"src_padding_mask\": BoolTensor[B, S],\n", " \"tgt_padding_mask\": BoolTensor[B, T],\n", "}\n", "```\n", "\n", "### Evaluation Interface (Person D -> C, E)\n", "```python\n", "evaluator = Evaluator(model, tokenizer, config)\n", "results = evaluator.evaluate(dataloader)\n", "# results: {\"bleu\": 25.6, \"comet\": 0.82, \"chrf\": 45.3, \"ter\": 55.2}\n", "```" ] } ], "metadata": { "colab": { "include_colab_link": true, "provenance": [] }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 4 }