优化ipynb

notebooks/EasyTranslate_Production.ipynb

@@ -37,6 +37,9 @@
     "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",
@@ -48,7 +51,7 @@
     "\n",
     "print(f\"Running in Google Colab: {IN_COLAB}\")\n",
     "print(f\"Python version: {sys.version}\")\n",
-    "print(f\"Working directory: {os.getcwd()}\")"
+    "print(f\"Working directory: {os.getcwd()}\")\n"
    ]
   },
   {
@@ -67,20 +70,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# You can override this in Colab before running this cell:\n",
-    "# os.environ[\"EASYTRANSLATE_REPO_URL\"] = \"https://github.com/<org>/<repo>.git\"\n",
+    "# ── 仓库地址配置 ─────────────────────────────────────────────────────────────\n",
+    "# 如需替换为你自己的仓库地址，在 Colab 中运行此单元前执行:\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",
-    "REPO_DIR = Path(\"/content/UCAS-EasyTranslate\") if IN_COLAB else Path(\".\").resolve()\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) -> Path | None:\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",
@@ -88,60 +92,59 @@
     "    return None\n",
     "\n",
     "\n",
-    "resolved_repo: Path | None = None\n",
+    "resolved_repo = None\n",
     "\n",
     "if IN_COLAB:\n",
-    "    # 1) If already present, use it directly\n",
-    "    if _is_repo_root(REPO_DIR):\n",
-    "        resolved_repo = REPO_DIR\n",
+    "    if _is_repo_root(_DEFAULT_COLAB_DIR):\n",
+    "        resolved_repo = _DEFAULT_COLAB_DIR\n",
     "    else:\n",
-    "        # 2) Try clone (public repo expected). If private/auth fails, continue to fallbacks.\n",
+    "        print(f\"Cloning repository from: {REPO_URL}\")\n",
     "        try:\n",
     "            subprocess.run(\n",
-    "                [\"git\", \"clone\", REPO_URL, str(REPO_DIR)],\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,\n",
-    "                stderr=subprocess.STDOUT,\n",
-    "                text=True,\n",
+    "                stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,\n",
     "            )\n",
-    "            resolved_repo = REPO_DIR\n",
-    "            print(f\"Cloned repository to: {REPO_DIR}\")\n",
+    "            resolved_repo = _DEFAULT_COLAB_DIR\n",
+    "            print(f\"Cloned to: {_DEFAULT_COLAB_DIR}\")\n",
     "        except subprocess.CalledProcessError as e:\n",
-    "            print(\"Clone failed. This usually means repo is private or URL is incorrect.\")\n",
-    "            print(\"Git output:\")\n",
-    "            print(e.stdout)\n",
+    "            print(f\"Clone failed:\\n{e.stdout}\")\n",
     "\n",
-    "    # 3) Fallback: try to locate an existing copy (e.g., from Drive or current working dir)\n",
+    "    # Fallback: search Drive or current dir\n",
     "    if resolved_repo is None:\n",
-    "        candidate_paths = [\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",
-    "        for candidate in candidate_paths:\n",
+    "        ]:\n",
     "            root = _find_repo_root(candidate)\n",
-    "            if root is not None:\n",
+    "            if root:\n",
     "                resolved_repo = root\n",
-    "                print(f\"Using existing repository at: {resolved_repo}\")\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",
-    "            \"Please do one of these: \"\n",
-    "            \"(1) set a public/correct EASYTRANSLATE_REPO_URL, \"\n",
-    "            \"(2) clone repo to /content/UCAS-EasyTranslate, \"\n",
-    "            \"(3) place repo in Google Drive and mount Drive.\"\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",
-    "    # Local: use nearest project root from current path\n",
     "    resolved_repo = _find_repo_root(Path.cwd()) or Path.cwd()\n",
     "    print(f\"Using local repository at: {resolved_repo}\")\n",
     "\n",
-    "os.chdir(resolved_repo)\n",
-    "sys.path.insert(0, str(Path(resolved_repo) / \"src\"))\n",
-    "print(f\"Repository directory: {resolved_repo}\")"
+    "# 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"
    ]
   },
   {
@@ -160,19 +163,33 @@
    "outputs": [],
    "source": [
     "if IN_COLAB:\n",
-    "    # Reuse Colab's preinstalled torch / numpy stack to avoid massive downgrades.\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",
-    "    %pip install -q -r requirements-colab.txt\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(\"Dependencies installed successfully\")\n",
-    "    print(\"Using Colab's preinstalled torch/numpy stack to minimize dependency conflicts.\")\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",
-    "\n",
-    "    print(\"Dependencies installed successfully\")\n"
+    "    print(\"Dependencies installed successfully.\")\n"
    ]
   },
   {
@@ -253,18 +270,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import importlib\n",
+    "import torch\n",
+    "import numpy as np\n",
     "\n",
-    "try:\n",
-    "    import numpy as np\n",
-    "    import pandas as pd\n",
-    "    print(f\"NumPy version: {np.__version__}\")\n",
-    "    print(f\"Pandas version: {pd.__version__}\")\n",
-    "except Exception as exc:\n",
-    "    raise RuntimeError(\n",
-    "        \"Numeric Python stack is broken in this runtime. \"\n",
-    "        \"Rerun the dependency installation cell, then restart the kernel and run the notebook from the top.\"\n",
-    "    ) from exc\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",
@@ -273,8 +283,8 @@
     "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",
+    "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",
@@ -283,19 +293,35 @@
     "    log_file=\"easytranslate.log\",\n",
     ")\n",
     "\n",
-    "# Colab overrides\n",
+    "# ── Colab runtime overrides ───────────────────────────────────────────────────\n",
     "if IN_COLAB:\n",
-    "    if not torch.cuda.is_available():\n",
-    "        config_dict[\"training\"][\"fp16\"] = False\n",
-    "        config_dict[\"training\"][\"bf16\"] = False\n",
-    "    config_dict[\"training\"][\"epochs\"] = config_dict[\"training\"].get(\"epochs\", 30)\n",
-    "\n",
-    "print(f\"Configuration loaded successfully\")\n",
-    "print(f\"Model type: {config_dict['model']['type']}\")\n",
-    "print(f\"Tokenizer type: {config_dict['tokenizer']['type']}\")\n",
-    "print(f\"Training epochs: {config_dict['training']['epochs']}\")\n",
-    "print(f\"FP16: {config_dict['training']['fp16']}, BF16: {config_dict['training']['bf16']}\")\n",
-    "print(f\"Random seed: {seed}\")\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"
    ]
   },
   {
@@ -323,15 +349,15 @@
     ")\n",
     "from torch.utils.data import DataLoader\n",
     "\n",
-    "data_cfg = config_dict.get(\"data\", {})\n",
+    "data_cfg   = config_dict.get(\"data\", {})\n",
     "preproc_cfg = data_cfg.get(\"preprocessing\", {})\n",
-    "loader_cfg = data_cfg.get(\"dataloader\", {})\n",
+    "loader_cfg  = data_cfg.get(\"dataloader\", {})\n",
     "\n",
-    "# Cap dataset size for Colab to avoid memory/time issues on large WMT corpora\n",
-    "MAX_TRAIN_SAMPLES = 200_000 if IN_COLAB else None   # set None to use full dataset\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...\")\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",
@@ -340,32 +366,37 @@
     ")\n",
     "\n",
     "train_raw = raw_dataset[\"train\"]\n",
-    "val_raw = raw_dataset.get(\"validation\", raw_dataset.get(\"dev\", train_raw))\n",
     "\n",
-    "print(f\"Raw training samples: {len(train_raw['src'])}\")\n",
-    "print(f\"Raw validation samples: {len(val_raw['src'])}\")\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 cap before preprocessing to save time\n",
-    "train_src_raw = train_raw[\"src\"][:MAX_TRAIN_SAMPLES] if MAX_TRAIN_SAMPLES else train_raw[\"src\"]\n",
-    "train_tgt_raw = train_raw[\"tgt\"][:MAX_TRAIN_SAMPLES] if MAX_TRAIN_SAMPLES else train_raw[\"tgt\"]\n",
-    "val_src_raw   = val_raw[\"src\"][:MAX_VAL_SAMPLES]   if MAX_VAL_SAMPLES   else val_raw[\"src\"]\n",
-    "val_tgt_raw   = val_raw[\"tgt\"][:MAX_VAL_SAMPLES]   if MAX_VAL_SAMPLES   else val_raw[\"tgt\"]\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(\"Preprocessing training data...\")\n",
+    "print(f\"\\nPreprocessing training data ({len(train_src_raw)} samples)...\")\n",
     "train_src, train_tgt = preprocess_pipeline(\n",
-    "    train_src_raw,\n",
-    "    train_tgt_raw,\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",
+    "print(f\"Preprocessed training samples  : {len(train_src)}\")\n",
     "\n",
-    "print(\"Preprocessing validation data...\")\n",
+    "print(f\"Preprocessing validation data ({len(val_src_raw)} samples)...\")\n",
     "val_src, val_tgt = preprocess_pipeline(\n",
-    "    val_src_raw,\n",
-    "    val_tgt_raw,\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",
@@ -440,19 +471,26 @@
     "    label_pad_token_id=-100,\n",
     ")\n",
     "\n",
-    "batch_size = loader_cfg.get(\"batch_size\", 32)\n",
-    "# Use 2 workers on Colab; 0 on CPU-only to avoid multiprocessing overhead\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",
-    "use_dynamic = loader_cfg.get(\"dynamic_batching\", True)\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...\")\n",
-    "    train_lengths = [\n",
-    "        max(len(tokenizer.encode(s, add_special_tokens=True)),\n",
-    "            len(tokenizer.encode(t, add_special_tokens=False)))\n",
-    "        for s, t in zip(train_src, train_tgt)\n",
-    "    ]\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",
@@ -466,6 +504,8 @@
     "        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",
@@ -484,11 +524,11 @@
     "    pin_memory=torch.cuda.is_available(),\n",
     ")\n",
     "\n",
-    "print(f\"Training batches: ~{len(train_loader)}\")\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(f\"Sample batch shapes:\")\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"
@@ -511,11 +551,25 @@
    "source": [
     "from easytranslate.model import TransformerTranslationModel\n",
     "\n",
-    "model_cfg = config_dict.get(\"model\", {})\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 = model_cfg.get(\"transformer\", {})\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",
@@ -533,7 +587,7 @@
     "        pad_id=tokenizer.pad_token_id,\n",
     "        share_embedding=False,\n",
     "    )\n",
-    "    print(f\"Built Transformer from scratch\")\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",
@@ -557,11 +611,11 @@
     "else:\n",
     "    raise ValueError(f\"Unknown model type: {model_type}\")\n",
     "\n",
-    "total_params = sum(p.numel() for p in model.parameters())\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}%\")"
+    "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"
    ]
   },
   {
@@ -616,6 +670,15 @@
    "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",
@@ -633,21 +696,21 @@
     "    evaluator=evaluator,\n",
     ")\n",
     "\n",
-    "# Unified output folders for artifacts and visualizations\n",
+    "# Output / plot directories live inside the repo\n",
     "OUTPUT_DIR = REPO_DIR / \"outputs\"\n",
-    "PLOTS_DIR = OUTPUT_DIR / \"plots\"\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: {trainer.fp16}, BF16: {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(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(\"To start training, run the next cell.\")"
+    "print(\">>> Run the NEXT cell to start training.\")\n"
    ]
   },
   {
@@ -692,21 +755,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json\n",
+    "EVAL_RESULTS = {}\n",
     "\n",
     "best_ckpt = trainer.checkpoint_dir / \"best_model.pt\"\n",
-    "EVAL_RESULTS = {}\n",
     "\n",
     "if best_ckpt.exists():\n",
     "    print(f\"Loading best model from {best_ckpt}\")\n",
-    "    checkpoint = torch.load(best_ckpt, map_location=device)\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...\")\n",
+    "    print(\"Running evaluation on validation set...\")\n",
     "    EVAL_RESULTS = evaluator.evaluate(\n",
     "        val_loader,\n",
     "        src_texts=val_src,\n",
@@ -717,15 +779,15 @@
     "    print(\"  Evaluation Results\")\n",
     "    print(\"=\" * 60)\n",
     "    for metric, score in EVAL_RESULTS.items():\n",
-    "        if not isinstance(score, list):\n",
-    "            print(f\"  {metric:>10s}: {score:.4f}\")\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\"Saved evaluation results to: {eval_path}\")\n",
+    "    print(f\"\\nSaved evaluation results → {eval_path}\")\n",
     "else:\n",
-    "    print(\"No best model checkpoint found. Run training first.\")"
+    "    print(\"No best_model.pt found. Run training first (Cell 13).\")\n"
    ]
   },
   {
@@ -743,12 +805,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json\n",
-    "\n",
     "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",
@@ -766,9 +827,9 @@
     "    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 to: {translation_path}\")\n",
+    "    print(f\"Saved translation examples → {translation_path}\")\n",
     "else:\n",
-    "    print(\"No trained model available for translation demo.\")"
+    "    print(\"No trained model checkpoint available. Run training first (Cell 13).\")\n"
    ]
   },
   {
@@ -786,14 +847,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import shutil\n",
     "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(trainer.log_dir),\n",
+    "        log_dir=str(_log_dir),\n",
     "        drive_base_path=DRIVE_BASE,\n",
     "    )\n",
     "\n",
@@ -807,24 +870,26 @@
     "    ]:\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",
-    "    # Sync plot images if they exist\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",
-    "    print(f\"Sync results: {sync_results}\")\n",
-    "    print(f\"Extra outputs synced to: {drive_outputs_dir}\")\n",
     "elif not IN_COLAB:\n",
-    "    print(f\"Running locally. Artifacts saved to:\")\n",
-    "    print(f\"  Checkpoints: {trainer.checkpoint_dir}\")\n",
-    "    print(f\"  Logs: {trainer.log_dir}\")\n",
-    "    print(f\"  Outputs: {OUTPUT_DIR}\")\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(\"Re-run with Drive mount to persist results.\")"
+    "    print(\"Mount Drive (Cell 5) and re-run this cell to sync results.\")\n"
    ]
   },
   {
@@ -842,10 +907,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json\n",
-    "\n",
-    "summary_path = trainer.checkpoint_dir / \"training_summary.json\"\n",
     "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",
@@ -854,27 +917,29 @@
     "    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')}): {TRAINING_SUMMARY.get('best_metric', 'N/A')}\")\n",
-    "    print(f\"  Total steps: {TRAINING_SUMMARY.get('total_steps', 'N/A')}\")\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",
-    "    if TRAINING_SUMMARY.get(\"train_loss_history\"):\n",
-    "        losses = TRAINING_SUMMARY[\"train_loss_history\"]\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\"  Loss reduction: {losses[0] - losses[-1]:.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 if \"EVAL_RESULTS\" in globals() else {},\n",
-    "        \"translation_examples\": TRANSLATION_RESULTS if \"TRANSLATION_RESULTS\" in globals() else [],\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\"Saved merged report to: {report_path}\")\n",
+    "    print(f\"\\nMerged report saved → {report_path}\")\n",
     "else:\n",
-    "    print(\"Training summary not yet available. Complete training first.\")"
+    "    print(\"Training summary not yet available. Complete training first.\")\n"
    ]
   },
   {