src/omnisub/config.py

"""Đọc và hợp nhất cấu hình từ `config.yaml`."""

from __future__ import annotations

import copy
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict

try:
    import yaml
except ImportError:  # pragma: no cover
    yaml = None  # type: ignore


DEFAULT_CONFIG: Dict[str, Any] = {
    "models": {
        "omni": "cyankiwi/Qwen3-Omni-30B-A3B-Instruct-AWQ-4bit",
        "quant": "awq",
        "diarize": "pyannote/speaker-diarization-community-1",
    },
    "translate": {
        "source_lang": "Chinese",
        "target_lang": "Vietnamese",
        "chars_per_sec": 22,
        "max_line_chars": 52,
        "max_lines": 2,
        "correct_ocr": False,
        "output_suffix": ".vi.srt",
    },
    "scene": {"max_gap": 1.5, "max_cues": 4, "max_dur": 20.0},
    "profiling": {"clips_per_speaker": 5, "max_seconds_per_clip": 8},
    "sampling": {
        "temperature": 0.7,
        "top_p": 0.95,
        "top_k": 64,
        "max_new_tokens": 1024,
    },
    "paths": {"drive_root": "/content/drive/MyDrive/Gemma"},
}


def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
    out = copy.deepcopy(base)
    for k, v in (override or {}).items():
        if isinstance(v, dict) and isinstance(out.get(k), dict):
            out[k] = _deep_merge(out[k], v)
        else:
            out[k] = v
    return out


@dataclass
class Config:
    data: Dict[str, Any]

    @classmethod
    def load(cls, path: str | Path | None = None) -> "Config":
        merged = copy.deepcopy(DEFAULT_CONFIG)
        if path is not None:
            p = Path(path)
            if p.exists():
                if yaml is None:
                    raise RuntimeError(
                        "Cần PyYAML để đọc config.yaml — `pip install pyyaml`."
                    )
                loaded = yaml.safe_load(p.read_text(encoding="utf-8")) or {}
                merged = _deep_merge(merged, loaded)
        return cls(data=merged)

    def __getitem__(self, key: str) -> Any:
        return self.data[key]

    def get(self, key: str, default: Any = None) -> Any:
        return self.data.get(key, default)

    @property
    def models(self) -> Dict[str, Any]:
        return self.data["models"]

    @property
    def translate(self) -> Dict[str, Any]:
        return self.data["translate"]

    @property
    def scene(self) -> Dict[str, Any]:
        return self.data["scene"]

    @property
    def profiling(self) -> Dict[str, Any]:
        return self.data["profiling"]

    @property
    def sampling(self) -> Dict[str, Any]:
        return self.data["sampling"]

    @property
    def paths(self) -> Dict[str, Any]:
        return self.data["paths"]