0fd2f875
 tangwang
translate
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
|
#!/usr/bin/env python3
"""Download local translation models declared in services.translation.capabilities."""
from __future__ import annotations
import argparse
from pathlib import Path
import os
import sys
from typing import Iterable
from huggingface_hub import snapshot_download
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
from config.services_config import get_translation_config
LOCAL_BACKENDS = {"local_nllb", "local_marian"}
def iter_local_capabilities(selected: set[str] | None = None) -> Iterable[tuple[str, dict]]:
cfg = get_translation_config()
for name, capability in cfg.capabilities.items():
backend = str(capability.get("backend") or "").strip().lower()
if backend not in LOCAL_BACKENDS:
continue
if selected and name not in selected:
continue
yield name, capability
def main() -> None:
parser = argparse.ArgumentParser(description="Download local translation models")
parser.add_argument("--all-local", action="store_true", help="Download all configured local translation models")
parser.add_argument("--models", nargs="*", default=[], help="Specific capability names to download")
args = parser.parse_args()
selected = {item.strip().lower() for item in args.models if item.strip()} or None
if not args.all_local and not selected:
parser.error("pass --all-local or --models <name> ...")
for name, capability in iter_local_capabilities(selected):
model_id = str(capability.get("model_id") or "").strip()
model_dir = Path(str(capability.get("model_dir") or "")).expanduser()
if not model_id or not model_dir:
raise ValueError(f"Capability '{name}' must define model_id and model_dir")
model_dir.parent.mkdir(parents=True, exist_ok=True)
print(f"[download] {name} -> {model_dir} ({model_id})")
snapshot_download(
repo_id=model_id,
local_dir=str(model_dir),
)
print(f"[done] {name}")
if __name__ == "__main__":
main()
|