asterocr 1.0.0__py3-none-any.whl

asterocr/kaggle_bridge/launcher.py

@@ -0,0 +1,245 @@
+"""
+AsterOCR Kaggle Launcher & Dispatcher
+Orchestrates Cloudflare Tunnel, prepares staging kernel script,
+pushes to Kaggle via official Kaggle CLI, and tracks remote execution lifecycle.
+"""
+
+import os
+import sys
+import time
+import json
+import uuid
+import threading
+import subprocess
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
+from asterocr.kaggle_bridge.tunnel import default_tunnel
+
+class KaggleLauncher:
+    def __init__(self, kernel_slug="creyt2012/asterocr-cloud-gpu-worker-v2"):
+        self.kernel_slug = kernel_slug
+        self.is_running = False
+        self.session_token = None
+        self.tunnel_url = None
+        self.monitor_thread = None
+        self.stop_requested = False
+
+    def _get_kaggle_env(self):
+        env = os.environ.copy()
+        tok_p = os.path.expanduser("~/.kaggle/access_token")
+        if os.path.exists(tok_p):
+            try:
+                env["KAGGLE_API_TOKEN"] = open(tok_p, "r", encoding="utf-8").read().strip()
+            except Exception:
+                pass
+        return env
+
+    def launch(self, config=None, on_status_update=None):
+        cfg = config or {}
+        self.stop_requested = False
+        self.session_token = f"sess_{uuid.uuid4().hex[:8]}"
+
+        # 1. Start Cloudflare Tunnel to expose local port 7860
+        print("[*] Đang khởi động Cloudflare Tunnel bảo mật...")
+        try:
+            self.tunnel_url = default_tunnel.start(timeout=30)
+        except Exception as e:
+            return False, f"Lỗi khởi động Cloudflare Tunnel: {e}"
+
+        # 2. Prepare staging directory
+        base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+        staging_dir = os.path.join(base_dir, "scratch", "kaggle_run")
+        os.makedirs(staging_dir, exist_ok=True)
+
+        # 3. Read template worker script and inject runtime config
+        worker_src_path = os.path.join(os.path.dirname(__file__), "kaggle_worker.py")
+        with open(worker_src_path, "r", encoding="utf-8") as f:
+            code = f.read()
+
+        batch_sz = min(96, int(cfg.get("batch_size", 96)))
+        epochs = int(cfg.get("epochs", 5))
+        
+        models_dir = os.path.abspath(os.path.join(base_dir, "models"))
+        has_pretrained = os.path.exists(os.path.join(models_dir, "asternet_best.pth"))
+        default_lr = 0.00005 if has_pretrained else 0.0005
+        lr = float(cfg.get("lr", default_lr))
+
+        is_tpu = (cfg.get("device") == "kaggle_tpu")
+        if is_tpu:
+            self.kernel_slug = "creyt2012/asterocr-cloud-tpu-worker"
+        else:
+            self.kernel_slug = "creyt2012/asterocr-cloud-gpu-worker-v2"
+
+        import re
+        code = re.sub(r'IS_TPU_DEVICE\s*=\s*(True|False)', f'IS_TPU_DEVICE = {is_tpu}', code)
+        code = re.sub(r'TUNNEL_URL\s*=\s*["\'][^"\']+["\']', f'TUNNEL_URL = "{self.tunnel_url}"', code)
+        code = re.sub(r'SESSION_TOKEN\s*=\s*["\'][^"\']+["\']', f'SESSION_TOKEN = "{self.session_token}"', code)
+        code = re.sub(r'BATCH_SIZE\s*=\s*\d+', f'BATCH_SIZE = {batch_sz}', code)
+        code = re.sub(r'LEARNING_RATE\s*=\s*[\d\.]+', f'LEARNING_RATE = {lr}', code)
+        code = re.sub(r'MAX_EPOCHS\s*=\s*\d+', f'MAX_EPOCHS = {epochs}', code)
+
+        staged_script = os.path.join(staging_dir, "kaggle_worker.py")
+        with open(staged_script, "w", encoding="utf-8") as f:
+            f.write(code)
+
+        # Generate Jupyter Notebook (.ipynb)
+        nb = {
+            "cells": [
+                {
+                    "cell_type": "markdown",
+                    "metadata": {},
+                    "source": [
+                        "# 🚀 AsterOCR Cloud GPU Worker v2\n",
+                        "Autonomous Deep Learning OCR Worker running on Kaggle GPU Cloud (Dual Tesla T4 / P100).\n",
+                        "- Real-time streaming chunks via Cloudflare Tunnel\n",
+                        "- FP16 Mixed Precision Training\n",
+                        "- Instant checkpoint synchronization"
+                    ]
+                },
+                {
+                    "cell_type": "code",
+                    "execution_count": None,
+                    "metadata": {},
+                    "outputs": [],
+                    "source": ["!nvidia-smi"]
+                },
+                {
+                    "cell_type": "code",
+                    "execution_count": None,
+                    "metadata": {},
+                    "outputs": [],
+                    "source": [
+                        "%%writefile kaggle_worker.py\n",
+                        code
+                    ]
+                },
+                {
+                    "cell_type": "code",
+                    "execution_count": None,
+                    "metadata": {},
+                    "outputs": [],
+                    "source": ["!python kaggle_worker.py"]
+                }
+            ],
+            "metadata": {
+                "accelerator": "GPU",
+                "kernelspec": {
+                    "display_name": "Python 3",
+                    "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.10.12"
+                }
+            },
+            "nbformat": 4,
+            "nbformat_minor": 5
+        }
+        nb_filename = "asterocr_cloud_gpu_worker_v2.ipynb"
+        with open(os.path.join(staging_dir, nb_filename), "w", encoding="utf-8") as f:
+            json.dump(nb, f, indent=2, ensure_ascii=False)
+
+        # 4. Generate kernel metadata
+        metadata = {
+            "id": self.kernel_slug,
+            "title": "AsterOCR Cloud TPU Worker" if is_tpu else "AsterOCR Cloud GPU Worker v2",
+            "code_file": nb_filename,
+            "language": "python",
+            "kernel_type": "notebook",
+            "is_private": True,
+            "enable_gpu": not is_tpu,
+            "enable_tpu": is_tpu,
+            "enable_internet": True
+        }
+        if not is_tpu:
+            metadata["machine_shape"] = "NvidiaTeslaT4"
+        meta_path = os.path.join(staging_dir, "kernel-metadata.json")
+        with open(meta_path, "w", encoding="utf-8") as f:
+            json.dump(metadata, f, indent=2)
+
+        # Bundle latest checkpoint into staging_dir for instant zero-latency recovery
+        import shutil
+        models_dir = os.path.abspath(os.path.join(base_dir, "models"))
+        best_p = os.path.join(models_dir, "asternet_best.pth")
+        last_p = os.path.join(models_dir, "asternet_last.pth")
+        chosen_pth = best_p if os.path.exists(best_p) else last_p
+        if os.path.exists(chosen_pth) and os.path.getsize(chosen_pth) > 1000:
+            shutil.copyfile(chosen_pth, os.path.join(staging_dir, "weights.pth"))
+            print(f"[+] Đã tự động đóng gói checkpoint '{os.path.basename(chosen_pth)}' ({os.path.getsize(chosen_pth)/(1024*1024):.1f} MB) vào Kernel!")
+
+        # 5. Push kernel to Cloud Compute
+        dev_str = "Cloud TPU (v3-8)" if is_tpu else "Cloud GPU (Dual NVIDIA Tesla T4 x2)"
+        print(f"[*] Đang kết nối và kích hoạt {dev_str} ({self.kernel_slug})...")
+        push_cmd = [sys.executable, "-m", "kaggle.cli", "kernels", "push", "-p", staging_dir]
+        k_env = self._get_kaggle_env()
+        try:
+            res = subprocess.run(push_cmd, env=k_env, capture_output=True, text=True, timeout=120)
+            print("[Cloud GPU Push]", res.stdout, res.stderr)
+            if res.returncode != 0 and "Kernel version" not in res.stdout:
+                default_tunnel.stop()
+                return False, f"Lỗi kích hoạt {dev_str}: {res.stderr or res.stdout}"
+        except Exception as e:
+            default_tunnel.stop()
+            return False, f"Lỗi gọi API {dev_str}: {e}"
+
+        self.is_running = True
+        print(f"[+] Đã kích hoạt thành công phiên làm việc trên {dev_str}!")
+
+        # 6. Start background status monitor
+        self.monitor_thread = threading.Thread(
+            target=self._monitor_loop,
+            args=(on_status_update, is_tpu),
+            daemon=True
+        )
+        self.monitor_thread.start()
+        return True, f"{dev_str} Trainer đã bắt đầu (Tunnel: {self.tunnel_url})"
+
+    def _monitor_loop(self, on_status_update, is_tpu):
+        time.sleep(25)
+        consecutive_errors = 0
+        dev_str = "Cloud TPU" if is_tpu else "Cloud GPU"
+        k_env = self._get_kaggle_env()
+        while self.is_running and not self.stop_requested:
+            try:
+                cmd = [sys.executable, "-m", "kaggle.cli", "kernels", "status", self.kernel_slug]
+                res = subprocess.run(cmd, env=k_env, capture_output=True, text=True, timeout=20)
+                out = res.stdout.strip().lower()
+                if "complete" in out:
+                    print(f"[+] {dev_str} Kernel đã hoàn thành 100%!")
+                    if on_status_update:
+                        on_status_update("completed")
+                    break
+                elif "error" in out or "failed" in out:
+                    consecutive_errors += 1
+                    if consecutive_errors >= 3:
+                        print(f"[!] {dev_str} Kernel báo lỗi: {res.stdout}")
+                        if on_status_update:
+                            on_status_update("error", res.stdout)
+                        break
+                elif "running" in out or "queued" in out:
+                    consecutive_errors = 0
+            except Exception:
+                pass
+            time.sleep(10)
+
+        self.is_running = False
+        default_tunnel.stop()
+
+    def stop(self):
+        self.stop_requested = True
+        self.is_running = False
+        default_tunnel.stop()
+
+default_launcher = KaggleLauncher()

asterocr/kaggle_bridge/tunnel.py

@@ -0,0 +1,190 @@
+"""
+AsterOCR Kaggle Tunnel Manager
+Automates zero-config HTTPS Cloudflare Quick Tunnel using engine/cloudflared.exe
+Exposes local Flask port (7860) to secure public internet so Kaggle P100 can stream data.
+"""
+
+import os
+import re
+import sys
+import time
+import subprocess
+import threading
+import requests
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
+        sys.stderr.reconfigure(encoding='utf-8', errors='replace')
+    except Exception:
+        pass
+
+class TunnelManager:
+    def __init__(self, local_port=7860, cloudflared_path=None):
+        self.local_port = local_port
+        if cloudflared_path is None:
+            base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+            self.cloudflared_path = os.path.join(base_dir, "engine", "cloudflared.exe")
+        else:
+            self.cloudflared_path = cloudflared_path
+            
+        self.process = None
+        self.tunnel_url = None
+        self._stop_event = threading.Event()
+        self._watchdog_thread = None
+
+    def start(self, timeout=30):
+        # 0. Check if user configured a persistent custom static URL (e.g. named tunnel on user's own domain)
+        config_path = os.path.join(os.path.dirname(__file__), "..", "tunnel_config.json")
+        custom_token = None
+        custom_fixed_url = None
+        if os.path.exists(config_path):
+            try:
+                import json
+                with open(config_path, "r", encoding="utf-8") as f:
+                    cfg = json.load(f)
+                    fixed_url = cfg.get("custom_url", "").strip()
+                    custom_token = cfg.get("tunnel_token", "").strip()
+                    if fixed_url.startswith("http://") or fixed_url.startswith("https://"):
+                        if "trycloudflare.com" not in fixed_url:
+                            custom_fixed_url = fixed_url.rstrip("/")
+            except Exception as e:
+                print(f"[*] Đọc cấu hình tunnel: {e}")
+
+        if not os.path.exists(self.cloudflared_path):
+            raise FileNotFoundError(f"Không tìm thấy cloudflared.exe tại: {self.cloudflared_path}")
+
+        # Terminate any stale cloudflared instances cleanly
+        self.stop()
+        if sys.platform == "win32":
+            try:
+                subprocess.run(["taskkill", "/f", "/im", "cloudflared.exe"], capture_output=True, timeout=5)
+            except Exception:
+                pass
+
+        self._stop_event.clear()
+
+        if custom_token:
+            print(f"[*] Kích hoạt Cloudflare Named Tunnel với Token cố định cho domain: {custom_fixed_url}...")
+            cmd = [
+                self.cloudflared_path,
+                "tunnel",
+                "run",
+                "--token", custom_token
+            ]
+        else:
+            cmd = [
+                self.cloudflared_path,
+                "tunnel",
+                "--url", f"http://127.0.0.1:{self.local_port}",
+                "--no-autoupdate"
+            ]
+
+        # Use creationflags on Windows to avoid opening popup console
+        creation_flags = 0
+        if sys.platform == "win32":
+            creation_flags = subprocess.CREATE_NO_WINDOW if hasattr(subprocess, "CREATE_NO_WINDOW") else 0x08000000
+
+        self.process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1,
+            creationflags=creation_flags
+        )
+
+        start_time = time.time()
+        extracted_url = None
+
+        if custom_token:
+            extracted_url = custom_fixed_url
+            # Wait briefly to confirm connector established
+            while time.time() - start_time < timeout:
+                line = self.process.stdout.readline()
+                if line:
+                    print(f"[Cloudflare] {line.strip()}")
+                    if "Registered tunnel connection" in line or "Connection" in line and "registered" in line:
+                        break
+                time.sleep(0.1)
+        else:
+            while time.time() - start_time < timeout:
+                line = self.process.stdout.readline()
+                if not line:
+                    time.sleep(0.2)
+                    continue
+                
+                # Match *.trycloudflare.com URL
+                match = re.search(r'https://[a-zA-Z0-9-]+\.trycloudflare\.com', line)
+                if match:
+                    extracted_url = match.group(0)
+                    break
+
+        if not extracted_url:
+            self.stop()
+            raise RuntimeError("Không thể tạo Cloudflare Tunnel trong thời gian quy định (Timeout).")
+
+        self.tunnel_url = extracted_url
+        self._save_last_url(self.tunnel_url)
+        print(f"[+] Cloudflare Tunnel kích hoạt thành công: {self.tunnel_url}")
+
+        # Start background watchdog to auto-restart tunnel if cloudflared unexpectedly dies
+        if self._watchdog_thread is None or not self._watchdog_thread.is_alive():
+            self._watchdog_thread = threading.Thread(target=self._watchdog_loop, daemon=True)
+            self._watchdog_thread.start()
+
+        return self.tunnel_url
+
+    def _watchdog_loop(self):
+        """Monitors cloudflared process; if it drops while active, resurrects it automatically."""
+        while not self._stop_event.is_set():
+            time.sleep(5)
+            if self._stop_event.is_set():
+                break
+            if self.tunnel_url and self.process and self.process.poll() is not None:
+                print("[!] Cảnh báo: Cloudflare Tunnel bị ngắt bất ngờ! Đang tự động hồi phục...")
+                try:
+                    self.start(timeout=30)
+                except Exception as e:
+                    print(f"[!] Tự động hồi phục tunnel thất bại: {e}")
+
+    def _save_last_url(self, url):
+        try:
+            p = os.path.join(os.path.dirname(__file__), "..", "last_tunnel_url.txt")
+            with open(p, "w", encoding="utf-8") as f:
+                f.write(url)
+        except Exception:
+            pass
+
+        # Auto-publish to Cloud GPU auto-discovery channel so the notebook reconnects 100% automatically
+        try:
+            requests.post("https://ntfy.sh/asterocr_tunnel_creyt2012", data=url.encode("utf-8"), timeout=6)
+            print(f"[+] Đã tự động đồng bộ Tunnel URL ({url}) lên kênh Cloud GPU (Auto-Discovery)!")
+        except Exception:
+            pass
+
+    def get_url(self):
+        return self.tunnel_url
+
+    def is_alive(self):
+        if self.process and self.process.poll() is None and self.tunnel_url:
+            return True
+        return False
+
+    def stop(self):
+        self._stop_event.set()
+        if self.process:
+            try:
+                self.process.terminate()
+                self.process.wait(timeout=3)
+            except Exception:
+                try:
+                    self.process.kill()
+                except Exception:
+                    pass
+            self.process = None
+        self.tunnel_url = None
+        print("[*] Cloudflare Tunnel đã đóng.")
+
+# Global singleton
+default_tunnel = TunnelManager()

asterocr/models/backbone.py

@@ -0,0 +1,102 @@
+"""
+AsterOCR Vision Encoder Backbone
+High-Frequency Asymmetric ResNet-SVTR for Vietnamese Document Recognition:
+- HighFrequencyStem: Dual-branch feature stem preserving micro-pixel diacritics
+- Asymmetric Spatial Downsampling (ASD): Compresses height (32 -> 1) while preserving horizontal sequence length (W/4)
+- SVTRLocalBlock: Depthwise stroke-grouping convolution with Squeeze-and-Excitation micro-channel attention
+- Hallucination-proof: 100% strictly bound to optical pixel evidence (Non-autoregressive)
+"""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class SqueezeExcitation(nn.Module):
+    """Channel attention mechanism to dynamically amplify diacritic-sensitive feature channels."""
+    def __init__(self, channels, reduction=16):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        mid = max(8, channels // reduction)
+        self.fc = nn.Sequential(
+            nn.Linear(channels, mid, bias=False),
+            nn.SiLU(inplace=False),
+            nn.Linear(mid, channels, bias=False),
+            nn.Sigmoid()
+        )
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.fc(self.avg_pool(x).view(b, c)).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+class HighFrequencyStem(nn.Module):
+    def __init__(self, in_channels=3, out_channels=64):
+        super().__init__()
+        self.conv_main = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.SiLU(inplace=False)
+        )
+        lap = torch.tensor([[0., -1., 0.], [-1., 4., -1.], [0., -1., 0.]]).view(1, 1, 3, 3)
+        self.register_buffer('hp_kernel', lap.repeat(in_channels, 1, 1, 1))
+        self.hp_proj = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_channels)
+        )
+        self.gamma = nn.Parameter(torch.tensor([0.15]))
+
+    def forward(self, x):
+        main_feat = self.conv_main(x)
+        with torch.no_grad():
+            k = self.hp_kernel.to(dtype=x.dtype, device=x.device)
+            hp = F.conv2d(x, k, padding=1, groups=x.shape[1])
+        return main_feat + self.gamma * self.hp_proj(hp)
+
+class SVTRLocalBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super().__init__()
+        self.pw1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.act1 = nn.SiLU(inplace=False)
+        self.dw = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, groups=out_channels, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.act2 = nn.SiLU(inplace=False)
+        self.se = SqueezeExcitation(out_channels, reduction=max(4, out_channels // 16))
+        self.pw2 = nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(out_channels)
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        res = self.shortcut(x)
+        out = self.act1(self.bn1(self.pw1(x)))
+        out = self.act2(self.bn2(self.dw(out)))
+        out = self.se(out)
+        out = self.bn3(self.pw2(out))
+        return F.silu(out + res)
+
+class AsterResNetBackbone(nn.Module):
+    def __init__(self, in_channels=3, out_dim=512):
+        super().__init__()
+        self.stem = HighFrequencyStem(in_channels, 64)
+        self.layer1 = nn.Sequential(SVTRLocalBlock(64, 64), SVTRLocalBlock(64, 128, stride=(2, 1))) # H: 32 -> 16
+        self.layer2 = nn.Sequential(SVTRLocalBlock(128, 128), SVTRLocalBlock(128, 256, stride=(2, 2))) # H: 16 -> 8, W/2
+        self.layer3 = nn.Sequential(SVTRLocalBlock(256, 256), SVTRLocalBlock(256, 512, stride=(2, 2))) # H: 8 -> 4, W/4
+        self.layer4 = nn.Sequential(
+            SVTRLocalBlock(512, 512),
+            nn.Conv2d(512, out_dim, kernel_size=(4, 1), stride=(4, 1), bias=False), # H: 4 -> 1, W/4
+            nn.BatchNorm2d(out_dim),
+            nn.SiLU(inplace=False)
+        )
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = x.squeeze(2)          # (B, C, W/4)
+        return x.permute(2, 0, 1) # (T, B, C)

asterocr/models/language.py

@@ -0,0 +1,104 @@
+"""
+AsterOCR Implicit Language Modeling Head (ILMH) & Gated Fusion Layer
+Inspired by ABINet (CVPR 2021) and PARSeq (ECCV 2022).
+Provides bidirectional linguistic reasoning and autonomous error correction
+for Vietnamese OCR without sacrificing real-time inference speed.
+"""
+
+import math
+import torch
+from torch import nn
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model=512, max_len=256, dropout=0.1):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(1) # Shape: (max_len, 1, d_model)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        # x shape: (T, B, d_model)
+        seq_len = x.size(0)
+        if seq_len > self.pe.size(0):
+            # Dynamic computation for extra-long sequences beyond max_len
+            position = torch.arange(0, seq_len, dtype=torch.float, device=x.device).unsqueeze(1)
+            d_model = x.size(-1)
+            div_term = torch.exp(torch.arange(0, d_model, 2, device=x.device).float() * (-math.log(10000.0) / d_model))
+            pe = torch.zeros(seq_len, d_model, device=x.device)
+            pe[:, 0::2] = torch.sin(position * div_term)
+            pe[:, 1::2] = torch.cos(position * div_term)
+            pe = pe.unsqueeze(1)
+            x = x + pe
+        else:
+            x = x + self.pe[:seq_len]
+        return self.dropout(x)
+
+class AsterLanguageModelHead(nn.Module):
+    """
+    Lightweight Bidirectional Implicit Language Modeling Head.
+    Operates on sequence features (T, B, hidden_dim) to model Vietnamese
+    orthographic, diacritical, and semantic transitions.
+    """
+    def __init__(self, d_model=512, nhead=4, num_layers=2, dim_feedforward=1024, dropout=0.1):
+        super(AsterLanguageModelHead, self).__init__()
+        self.pos_encoder = PositionalEncoding(d_model=d_model, max_len=256, dropout=dropout)
+        
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation="gelu",
+            norm_first=True, # Pre-LayerNorm for stable gradient flow
+            batch_first=False
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        # x shape: (T, B, d_model)
+        pos_x = self.pos_encoder(x)
+        lm_feat = self.transformer(pos_x)
+        return self.norm(lm_feat)
+
+class GatedLanguageFusion(nn.Module):
+    """
+    Adaptive Gated Vision-Language Fusion:
+    Computes a dynamic gate vector G = sigmoid(Linear([H_vis; H_lang])) in (0, 1).
+    When vision is clear: G -> 1 (prioritize visual features).
+    When vision is blurred/degraded/ambiguous: G -> 0 (rely on LM to correct errors).
+    """
+    def __init__(self, d_model=512, init_vision_bias=1.5):
+        super(GatedLanguageFusion, self).__init__()
+        self.gate_linear = nn.Linear(d_model * 2, d_model)
+        # Initialize bias positively so initially G ~= sigmoid(1.5) ~= 0.82 (safe for existing weights)
+        nn.init.constant_(self.gate_linear.bias, init_vision_bias)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, h_vis, h_lang):
+        # h_vis, h_lang: (T, B, d_model)
+        concat = torch.cat([h_vis, h_lang], dim=-1)
+        gate = self.sigmoid(self.gate_linear(concat))
+        h_fused = gate * h_vis + (1.0 - gate) * h_lang
+        return h_fused, gate
+
+if __name__ == "__main__":
+    lm = AsterLanguageModelHead(d_model=512, nhead=4, num_layers=2)
+    fusion = GatedLanguageFusion(d_model=512)
+    
+    dummy_vis = torch.randn(64, 4, 512)
+    dummy_lang = lm(dummy_vis)
+    fused, gate = fusion(dummy_vis, dummy_lang)
+    
+    print("LM output shape:", dummy_lang.shape)
+    print("Fused output shape:", fused.shape)
+    print("Mean gate value (initially vision-biased):", gate.mean().item())
+    total_params = sum(p.numel() for p in lm.parameters()) + sum(p.numel() for p in fusion.parameters())
+    print(f"Total Language + Fusion Parameters: {total_params / 1e6:.3f}M")