{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Gemma SRT Translate — Google Colab\n", "\n", "Dịch phụ đề SRT bằng **Gemma 4 12B vision** (llama-server + MTP).\n", "\n", "**Yêu cầu:** GPU **T4/L4** — Runtime → Change runtime type → GPU.\n", "\n", "**Lần đầu:** tải model ~12GB (+ build llama-server một lần, lưu `llama-server-bin.tgz` lên Drive). \n", "**Các lần sau:** restore **`Gemma/Cache/llama-server-bin.tgz`** (~vài giây) — **không còn cell build dài**.\n", "\n", "---\n", "\n", "## Pipeline (2 bước)\n", "\n", "| Bước | Tên | Mặc định | Việc làm |\n", "|------|-----|----------|----------|\n", "| **Pass 1** | OCR / sửa SRT gốc | **Tắt** (`SKIP_CORRECTION = True`) | Cắt khung hình video → đọc sub burn-in → sửa lỗi chính tả trong `.srt` gốc |\n", "| **Pass 2** | Dịch | **Luôn chạy** | Dịch sang tiếng Việt (có ngữ cảnh scene + xưng hô) |\n", "\n", "**`SKIP_CORRECTION` — đọc nhanh:**\n", "- `True` (**mặc định**) → **bỏ Pass 1**, chỉ dịch → **nhanh hơn**\n", "- `False` → **chạy Pass 1** OCR/sửa SRT trước khi dịch → **chậm hơn**, dùng khi sub gốc hay sai chữ\n", "\n", "**Dịch tự nhiên (lồng tiếng):** budget rộng (22 ký tự/s, tối đa ~104 ký tự/cue), **tắt Pass 2b rút gọn** mặc định.\n", "\n", "---\n", "\n", "## Cách chạy\n", "\n", "1. Đặt phim + SRT vào **Drive → Gemma → Phim** (cùng tên, ví dụ `Phim A.mp4` + `Phim A.srt`)\n", "2. Sửa **`PHIM_STEMS`** ở cell cấu hình — tên file **không đuôi**, khớp y hệt trên Drive\n", "3. Mở từ [HF /colab](https://huggingface.co/STBack23/gemma-srt-translate/colab) → **Runtime → Run all**\n", "4. File `*.vi.srt` ra **Gemma/Phim** (model cache: **Gemma/Cache**)\n", "\n", "Sau khi xong: `DOWNLOAD_RESULTS` tải file về trình duyệt; `AUTO_UNMOUNT_DRIVE` / `AUTO_DISCONNECT_RUNTIME` tự gỡ Drive và ngắt runtime." ] }, { "cell_type": "code", "metadata": {}, "source": [ "# ═══ CẤU HÌNH — sửa trước khi chạy ═══\n", "from pathlib import Path\n", "\n", "HF_REPO = \"STBack23/gemma-srt-translate\"\n", "MODELS_REPO = \"unsloth/gemma-4-12B-it-qat-GGUF\" # repo model GGUF (~12GB) — không cần sửa\n", "\n", "# ── Google Drive ─────────────────────────────────────────────\n", "USE_DRIVE_CACHE = True # True = cache model + llama-server trên Drive (khuyến nghị)\n", "\n", "# Thư mục trên Drive: Gemma/Cache (model) + Gemma/Phim (video + SRT vào/ra)\n", "DRIVE_ROOT = \"/content/drive/MyDrive/Gemma\"\n", "DRIVE_CACHE = f\"{DRIVE_ROOT}/Cache\"\n", "DRIVE_PHIM_DIR = f\"{DRIVE_ROOT}/Phim\"\n", "\n", "VIDEO_EXTS = (\".mp4\", \".mkv\", \".avi\", \".mov\")\n", "\n", "def job(stem, ext_video=\".mp4\"):\n", " \"\"\"stem = tên file trên Drive (KHÔNG đuôi), khớp y hệt — kể cả dấu, khoảng trắng.\"\"\"\n", " base = f\"{DRIVE_PHIM_DIR}/{stem}\"\n", " return {\"stem\": stem, \"video\": f\"{base}{ext_video}\", \"srt\": f\"{base}.srt\"}\n", "\n", "# ── Danh sách phim cần dịch ──────────────────────────────────\n", "# Tên trong PHIM_STEMS phải trùng file trong Gemma/Phim (không gồm .mp4 / .srt)\n", "PHIM_STEMS = [\n", " \"Thư Gửi Mùa Hè\", # → Gemma/Phim/Thư Gửi Mùa Hè.mp4 + Thư Gửi Mùa Hè.srt\n", " # \"ten-phim-khac\",\n", "]\n", "JOBS = [job(stem) for stem in PHIM_STEMS]\n", "# Cách khác — đường dẫn đầy đủ:\n", "# JOBS = [{\"video\": f\"{DRIVE_PHIM_DIR}/a.mp4\", \"srt\": f\"{DRIVE_PHIM_DIR}/a.srt\"}]\n", "# Upload 1 phim qua nút Colab: PHIM_STEMS = [], JOBS = [], UPLOAD_WIDGET = True\n", "UPLOAD_WIDGET = False\n", "\n", "# ── Tùy chọn dịch ───────────────────────────────────────────\n", "SOURCE_LANG = \"auto\" # ngôn ngữ SRT gốc (\"auto\" = model tự nhận)\n", "TARGET_LANG = \"Vietnamese\" # ngôn ngữ đích\n", "\n", "# Pass 1 — OCR / sửa SRT gốc từ khung hình video (xem bảng ở đầu notebook):\n", "# True → BỎ QUA Pass 1 — chỉ dịch Pass 2 (MẶC ĐỊNH, nhanh)\n", "# False → CHẠY Pass 1 OCR/sửa sub trước — chậm hơn, dùng khi sub hay sai chữ\n", "SKIP_CORRECTION = True\n", "\n", "LIMIT_CUES = 0 # 0 = dịch hết file; đặt 20 để thử nhanh vài cue đầu\n", "\n", "# ── Sau khi dịch xong ─────────────────────────────────────────\n", "DOWNLOAD_RESULTS = True # True = tải *.vi.srt về trình duyệt\n", "AUTO_UNMOUNT_DRIVE = True # True = gỡ mount Google Drive\n", "AUTO_DISCONNECT_RUNTIME = True # True = ngắt runtime Colab (tiết kiệm GPU quota)\n", "DISCONNECT_DELAY_SEC = 15 # giây chờ trước khi ngắt (để download kịp)\n", "\n", "# ── Đường dẫn nội bộ Colab (thường không cần sửa) ────────────\n", "ROOT = \"/content/gemma-srt\"\n", "MODELS_DIR = f\"{ROOT}/models\"\n", "LLAMA_DIR = \"/content/llama.cpp\"\n", "LLAMA_SERVER = f\"{LLAMA_DIR}/build/bin/llama-server\"" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Kiểm tra GPU\n", "!nvidia-smi --query-gpu=name,memory.total --format=csv,noheader\n", "\n", "import torch\n", "if not torch.cuda.is_available():\n", " raise RuntimeError(\"Chưa có GPU. Runtime → Change runtime type → GPU (L4).\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Mount Google Drive — Gemma/Cache (model) + Gemma/Phim (phim)\n", "if USE_DRIVE_CACHE:\n", " from google.colab import drive\n", " from pathlib import Path\n", " import os\n", " drive.mount(\"/content/drive\")\n", " os.makedirs(DRIVE_CACHE, exist_ok=True)\n", " os.makedirs(DRIVE_PHIM_DIR, exist_ok=True)\n", " print(f\"Cache model: {DRIVE_CACHE}\")\n", " print(f\"Thư mục phim: {DRIVE_PHIM_DIR}\")\n", " cache_tgz = Path(f\"{DRIVE_CACHE}/llama-server-bin.tgz\")\n", " if cache_tgz.is_file():\n", " sz = cache_tgz.stat().st_size\n", " mb = sz / 1_048_576\n", " print(f\" llama-server-bin.tgz: {mb:.1f} MB\" + (\" OK\" if sz > 5_000_000 else \" HỎNG\"))\n", " else:\n", " for old in (\"llama-server.tgz\", \"llama-server\"):\n", " p = Path(f\"{DRIVE_CACHE}/{old}\")\n", " if p.is_file():\n", " print(f\" cache cũ {old}: {p.stat().st_size/1024:.0f} KB — sẽ thay bằng llama-server-bin.tgz\")\n", " if not any(Path(f\"{DRIVE_CACHE}/{n}\").is_file() for n in (\"llama-server-bin.tgz\", \"llama-server.tgz\", \"llama-server\")):\n", " print(\" llama-server cache: chưa có (lần đầu build ~20–50 phút)\")\n", " phim = Path(DRIVE_PHIM_DIR)\n", " files = sorted(p.name for p in phim.iterdir() if p.is_file())\n", " if files:\n", " print(f\" ({len(files)} file trong Phim)\")\n", " for name in files[:12]:\n", " print(f\" - {name}\")\n", " if len(files) > 12:\n", " print(f\" ... và {len(files) - 12} file khác\")\n", " else:\n", " print(\" (chưa có file — upload video + .srt vào Gemma/Phim)\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Cài dependency + tải code từ Hugging Face\n", "!pip install -q huggingface_hub pyyaml\n", "\n", "from huggingface_hub import snapshot_download\n", "from pathlib import Path\n", "import os\n", "\n", "if HF_REPO.startswith(\"YOUR_\"):\n", " raise ValueError(\"Sửa HF_REPO thành repo Hugging Face của bạn, ví dụ: 'username/gemma-srt-translate'\")\n", "\n", "print(f\"[code] Downloading {HF_REPO} ...\")\n", "snapshot_download(\n", " repo_id=HF_REPO,\n", " repo_type=\"model\",\n", " local_dir=ROOT,\n", " local_dir_use_symlinks=False,\n", ")\n", "print(f\"[code] OK: {ROOT}\")\n", "assert Path(f\"{ROOT}/translate_srt.py\").is_file(), \"translate_srt.py not found in repo\"" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Tải model GGUF (~12GB) — cache trên Drive nếu bật\n", "from huggingface_hub import hf_hub_download\n", "from pathlib import Path\n", "import os\n", "import shutil\n", "\n", "MODEL_FILES = [\n", " \"gemma-4-12B-it-qat-UD-Q4_K_XL.gguf\",\n", " \"mmproj-F16.gguf\",\n", " \"mtp-gemma-4-12B-it.gguf\",\n", "]\n", "\n", "models_dest = MODELS_DIR\n", "if USE_DRIVE_CACHE:\n", " models_dest = f\"{DRIVE_CACHE}/models\"\n", "Path(models_dest).mkdir(parents=True, exist_ok=True)\n", "\n", "paths = {}\n", "for name in MODEL_FILES:\n", " dest_file = Path(models_dest) / name\n", " if dest_file.is_file():\n", " print(f\"[model] cached: {name}\")\n", " paths[name] = dest_file\n", " continue\n", " print(f\"[model] downloading {MODELS_REPO}/{name} ...\")\n", " p = hf_hub_download(\n", " repo_id=MODELS_REPO,\n", " filename=name,\n", " local_dir=models_dest,\n", " local_dir_use_symlinks=False,\n", " )\n", " paths[name] = Path(p)\n", " print(f\" OK: {p}\")\n", "\n", "# Symlink/copy vào ROOT/models cho translate_srt.py\n", "Path(MODELS_DIR).mkdir(parents=True, exist_ok=True)\n", "for name, src in paths.items():\n", " dst = Path(MODELS_DIR) / name\n", " if not dst.exists():\n", " try:\n", " os.symlink(src, dst)\n", " except OSError:\n", " shutil.copy2(src, dst)\n", "\n", "MODEL_PATH = Path(MODELS_DIR) / MODEL_FILES[0]\n", "MMPROJ_PATH = Path(MODELS_DIR) / MODEL_FILES[1]\n", "DRAFT_PATH = Path(MODELS_DIR) / MODEL_FILES[2]\n", "print(f\"\\nModels ready in {MODELS_DIR}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# llama-server — restore từ Drive (Gemma/Cache/llama-server-bin.tgz)\n", "import sys\n", "\n", "sys.path.insert(0, ROOT)\n", "from scripts.ensure_llama_colab import ensure_llama_server\n", "\n", "LLAMA_SERVER = str(ensure_llama_server(\n", " llama_dir=LLAMA_DIR,\n", " drive_cache=DRIVE_CACHE if USE_DRIVE_CACHE else None,\n", " allow_build=False,\n", "))\n", "print(f\"LLAMA_SERVER = {LLAMA_SERVER}\")\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Kiểm tra ffmpeg\n", "!ffmpeg -version | head -1" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Kiểm tra danh sách JOBS (Drive) hoặc upload 1 phim qua widget\n", "from pathlib import Path\n", "from google.colab import files\n", "import os\n", "\n", "UPLOAD_DIR = \"/content/uploads\"\n", "os.makedirs(UPLOAD_DIR, exist_ok=True)\n", "resolved_jobs = []\n", "\n", "def _list_phim_hint() -> str:\n", " phim = Path(DRIVE_PHIM_DIR)\n", " if not phim.is_dir():\n", " return f\"Chưa có thư mục: {DRIVE_PHIM_DIR}\"\n", " lines = [\"File hiện có trong Gemma/Phim:\"]\n", " for p in sorted(phim.iterdir()):\n", " if p.is_file():\n", " lines.append(f\" - {p.name}\")\n", " if len(lines) == 1:\n", " lines.append(\" (trống — upload video + .srt vào đây)\")\n", " return \"\\n\".join(lines)\n", "\n", "def _find_video(stem: str, preferred: str | None = None) -> Path | None:\n", " if preferred:\n", " p = Path(preferred)\n", " if p.is_file():\n", " return p\n", " for ext in VIDEO_EXTS:\n", " p = Path(DRIVE_PHIM_DIR) / f\"{stem}{ext}\"\n", " if p.is_file():\n", " return p\n", " return None\n", "\n", "if JOBS:\n", " for i, entry in enumerate(JOBS, 1):\n", " stem = entry.get(\"stem\") or Path(entry[\"video\"]).stem\n", " v = _find_video(stem, entry.get(\"video\"))\n", " s = Path(entry[\"srt\"])\n", " if v is None:\n", " raise FileNotFoundError(\n", " f\"Job {i}: không thấy video cho '{stem}'\\n\"\n", " f\"Đã thử: {', '.join(stem + ext for ext in VIDEO_EXTS)}\\n\\n\"\n", " f\"Tên trong PHIM_STEMS phải KHỚP Y HỆT tên file (không đuôi).\\n\\n\"\n", " f\"{_list_phim_hint()}\"\n", " )\n", " if not s.is_file():\n", " s = Path(DRIVE_PHIM_DIR) / f\"{stem}.srt\"\n", " if not s.is_file():\n", " raise FileNotFoundError(\n", " f\"Job {i}: không thấy SRT: {s}\\n\\n\"\n", " f\"Cần file: {stem}.srt cùng thư mục Phim.\\n\\n\"\n", " f\"{_list_phim_hint()}\"\n", " )\n", " out = Path(entry[\"output\"]) if entry.get(\"output\") else v.with_name(v.stem + \".vi.srt\")\n", " resolved_jobs.append({\"video\": v, \"srt\": s, \"output\": out})\n", " print(f\"Job {i}/{len(JOBS)}: {v.name} + {s.name} -> {out.name}\")\n", "elif UPLOAD_WIDGET:\n", " print(\"Upload 1 video (.mp4/.mkv) và 1 SRT (.srt):\")\n", " uploaded = files.upload()\n", " video = srt = None\n", " for name in uploaded:\n", " p = Path(UPLOAD_DIR) / name\n", " p.write_bytes(uploaded[name])\n", " low = name.lower()\n", " if low.endswith((\".mp4\", \".mkv\", \".avi\", \".mov\")):\n", " video = p\n", " elif low.endswith(\".srt\"):\n", " srt = p\n", " if not video or not srt:\n", " raise ValueError(\"Cần 1 file video và 1 file .srt\")\n", " out = video.with_name(video.stem + \".vi.srt\")\n", " resolved_jobs.append({\"video\": video, \"srt\": srt, \"output\": out})\n", " print(f\"Job 1/1: {video.name} -> {out.name}\")\n", "else:\n", " raise ValueError(\"Thêm phim vào JOBS hoặc đặt UPLOAD_WIDGET = True\")\n", "\n", "print(f\"\\nTổng: {len(resolved_jobs)} phim (chạy tuần tự)\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# ═══ CHẠY DỊCH SRT (tuần tự từng phim) ═══\n", "import sys\n", "import os\n", "from pathlib import Path\n", "\n", "sys.path.insert(0, ROOT)\n", "from scripts.ensure_llama_colab import ensure_llama_server, llama_bin_valid\n", "\n", "if not llama_bin_valid(Path(LLAMA_SERVER)):\n", " LLAMA_SERVER = str(ensure_llama_server(\n", " llama_dir=LLAMA_DIR,\n", " drive_cache=DRIVE_CACHE if USE_DRIVE_CACHE else None,\n", " allow_build=False,\n", " ))\n", "\n", "from translate_srt import build_parser, run_pipeline\n", "\n", "completed = []\n", "\n", "for i, job in enumerate(resolved_jobs, 1):\n", " v, s, out = job[\"video\"], job[\"srt\"], job[\"output\"]\n", " print(\"\\n\" + \"=\" * 50)\n", " print(f\"=== Phim {i}/{len(resolved_jobs)}: {v.name} ===\")\n", " print(\"=\" * 50)\n", "\n", " argv = [\n", " \"--video\", str(v),\n", " \"--input-srt\", str(s),\n", " \"--output-srt\", str(out),\n", " \"--source-lang\", SOURCE_LANG,\n", " \"--target-lang\", TARGET_LANG,\n", " \"--llama-server\", LLAMA_SERVER,\n", " \"--model\", str(MODEL_PATH),\n", " \"--mmproj\", str(MMPROJ_PATH),\n", " \"--model-draft\", str(DRAFT_PATH),\n", " \"--limit\", str(LIMIT_CUES),\n", " \"--ngl\", \"999\",\n", " \"--ctx\", \"8192\",\n", " ]\n", " if SKIP_CORRECTION:\n", " argv.append(\"--skip-correction\")\n", "\n", " args = build_parser().parse_args(argv)\n", " args.log_fn = lambda msg, _i=i: print(f\"[{_i}] {msg}\", flush=True)\n", " run_pipeline(args)\n", " completed.append(out)\n", " print(f\"\\nXong phim {i}: {out}\")\n", "\n", "print(\"\\n\" + \"=\" * 50)\n", "print(f\"HOÀN TẤT {len(completed)} phim:\")\n", "for p in completed:\n", " print(f\" - {p}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "# Tải kết quả về máy (tuỳ chọn) + ngắt kết nối\n", "from pathlib import Path\n", "import time\n", "\n", "if DOWNLOAD_RESULTS:\n", " from google.colab import files\n", " for out in completed:\n", " out = Path(out)\n", " for p in [\n", " out,\n", " out.with_suffix(\".corrected\" + out.suffix),\n", " out.with_suffix(out.suffix + \".report.json\"),\n", " ]:\n", " if p.is_file():\n", " print(f\"Download: {p.name}\")\n", " files.download(str(p))\n", "else:\n", " print(\"Bỏ qua download — file SRT đã nằm trên Drive (xem đường dẫn output trong JOBS).\")\n", "\n", "if AUTO_UNMOUNT_DRIVE and USE_DRIVE_CACHE:\n", " from google.colab import drive\n", " drive.flush_and_unmount()\n", " print(\"Đã gỡ mount Google Drive.\")\n", "\n", "if AUTO_DISCONNECT_RUNTIME:\n", " delay = max(0, int(DISCONNECT_DELAY_SEC))\n", " if delay:\n", " print(f\"Ngắt Colab runtime sau {delay}s...\")\n", " time.sleep(delay)\n", " from google.colab import runtime\n", " runtime.unassign()" ], "execution_count": null, "outputs": [] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "L4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }