offline-translate 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,11 @@
1
+ .venv/
2
+ __pycache__/
3
+ *.egg-info/
4
+ dist/
5
+ .pytest_cache/
6
+ .DS_Store
7
+
8
+ # downloaded language packs and the phone dump (contains personal data)
9
+ models/
10
+ phone_dump/
11
+ .claude/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 asimon655
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,136 @@
1
+ Metadata-Version: 2.5
2
+ Name: offline-translate
3
+ Version: 0.1.0
4
+ Summary: Offline machine translation using Google Translate's on-device language packs
5
+ Project-URL: Homepage, https://github.com/asimon655/offline-translate
6
+ Project-URL: Issues, https://github.com/asimon655/offline-translate/issues
7
+ Author: asimon655
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: google-translate,nmt,offline,tflite,translation
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Operating System :: MacOS
14
+ Classifier: Operating System :: Microsoft :: Windows
15
+ Classifier: Operating System :: POSIX :: Linux
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3 :: Only
18
+ Classifier: Topic :: Text Processing :: Linguistic
19
+ Classifier: Typing :: Typed
20
+ Requires-Python: >=3.10
21
+ Requires-Dist: ai-edge-litert>=1.0
22
+ Requires-Dist: numpy>=1.24
23
+ Requires-Dist: sentencepiece>=0.2
24
+ Provides-Extra: test
25
+ Requires-Dist: pytest>=8; extra == 'test'
26
+ Description-Content-Type: text/markdown
27
+
28
+ # offline-translate
29
+
30
+ Offline machine translation in Python using the same on-device neural translation models that the
31
+ Google Translate Android app downloads for offline use.
32
+
33
+ - ~100 languages, each pack translates to/from English (~50 MB per language)
34
+ - Fully offline once a pack is downloaded; no API key, no account
35
+ - Language pairs without English (e.g. Hebrew → French) pivot through English, like the app does
36
+ - Runs on CPU via [LiteRT](https://ai.google.dev/edge/litert) (TensorFlow Lite); ~0.3 s per sentence
37
+
38
+ ```bash
39
+ pip install offline-translate
40
+ ```
41
+
42
+ Supported platforms are those with `ai-edge-litert` wheels: macOS (Apple Silicon), Linux (x86_64,
43
+ aarch64) and Windows (x64), Python 3.10+.
44
+
45
+ ## Usage
46
+
47
+ ```python
48
+ from offline_translate import Translator
49
+
50
+ t = Translator("~/.local/share/offline-translate") # where models are stored
51
+
52
+ t.download("he") # Hebrew <-> English, ~55 MB
53
+ t.translate("Good morning!", target="he").text
54
+ # 'בוקר טוב!'
55
+
56
+ t.translate("בוקר טוב", target="en") # source is auto-detected among installed languages
57
+ # Translation(text='Good morning', source='iw', target='en', via=('iw', 'en'))
58
+ ```
59
+
60
+ ### Managing languages
61
+
62
+ ```python
63
+ for lang in t.languages(): # full catalog (fetched once from Google, then cached)
64
+ print(lang.code, lang.name, lang.installed, lang.download_size, lang.update_available)
65
+
66
+ t.installed_languages() # ['en', 'iw'] — works offline, no catalog needed
67
+ t.download("fr", progress=lambda done, total: print(f"{done}/{total}"))
68
+ t.remove("fr")
69
+
70
+ t.update_catalog() # check Google for a newer catalog
71
+ [l for l in t.languages() if l.update_available] # then t.download(code) to update
72
+ ```
73
+
74
+ Language codes follow Google's catalog (`iw` for Hebrew, `zh-Hant`, `jw`, …). Common aliases
75
+ (`he`, `zh-TW`, `jv`) and English names (`"Hebrew"`) are accepted everywhere.
76
+
77
+ English is built in: there is no separate English pack, since every pack contains both
78
+ directions to and from English.
79
+
80
+ ### API
81
+
82
+ | | |
83
+ |---|---|
84
+ | `Translator(models_dir, channel="prod", beam=4)` | `models_dir` is created if missing. `beam=1` is faster, slightly worse. |
85
+ | `translate(text, target, source=None) -> Translation` | `Translation(text, source, target, via)` |
86
+ | `detect(text) -> str` | Guess among installed languages |
87
+ | `languages(installed_only=False) -> list[Language]` | `Language(code, name, installed, download_size, size, update_available, builtin, supported)` |
88
+ | `installed_languages() -> list[str]` | |
89
+ | `download(lang, progress=None, force=False)` | Verifies SHA-1 of every file; resumes interrupted downloads |
90
+ | `remove(lang)` | |
91
+ | `update_catalog() -> str` | Returns the catalog version |
92
+ | `TranslateError` | Raised for unknown/uninstalled languages, checksum failures, … |
93
+
94
+ Network errors surface as `OSError`.
95
+
96
+ ## Command line
97
+
98
+ The repository also has a small CLI built on the library (not part of the PyPI package):
99
+
100
+ ```bash
101
+ ./translate get he # download Hebrew
102
+ ./translate "Good morning" # English -> Hebrew by default; source auto-detected
103
+ ./translate -s he -t fr "בוקר טוב" # any pair, via English
104
+ ./translate list --installed
105
+ ```
106
+
107
+ `tools/pull_from_phone.py` copies already-downloaded packs from a rooted Android phone over adb.
108
+
109
+ ## How it works
110
+
111
+ The Android app reads a manifest at
112
+ `https://dl.google.com/translate/offline/manifest/gtr/prod/manifest.pb`, which lists versions of
113
+ a *package profile* — the catalog of language packs and their files. Each translation direction is
114
+ a TFLite encoder, an LSTM decoder run step by step with beam search, and SentencePiece tokenizers.
115
+
116
+ ### Limitations
117
+
118
+ - 15 languages (am, az, bs, hy, kk, ky, ml, mn, ne, pa, ps, si, so, tg, uz) still ship an older
119
+ model format with a WordPiece tokenizer that isn't supported yet; they're listed with
120
+ `supported=False`.
121
+ - Quality matches the app's offline mode, which is below online Google Translate.
122
+
123
+ ## Development
124
+
125
+ ```bash
126
+ pip install -e ".[test]"
127
+ pytest # offline unit tests (fake Google server, fake model)
128
+ pytest --network # + checks against the real Google catalog
129
+ pytest --models=~/models # + end-to-end translation with installed Hebrew models
130
+ ```
131
+
132
+ ## Disclaimer
133
+
134
+ This project is not affiliated with or endorsed by Google. It contains no Google code or models;
135
+ language packs are downloaded from Google's public servers at runtime and are subject to Google's
136
+ terms.
@@ -0,0 +1,109 @@
1
+ # offline-translate
2
+
3
+ Offline machine translation in Python using the same on-device neural translation models that the
4
+ Google Translate Android app downloads for offline use.
5
+
6
+ - ~100 languages, each pack translates to/from English (~50 MB per language)
7
+ - Fully offline once a pack is downloaded; no API key, no account
8
+ - Language pairs without English (e.g. Hebrew → French) pivot through English, like the app does
9
+ - Runs on CPU via [LiteRT](https://ai.google.dev/edge/litert) (TensorFlow Lite); ~0.3 s per sentence
10
+
11
+ ```bash
12
+ pip install offline-translate
13
+ ```
14
+
15
+ Supported platforms are those with `ai-edge-litert` wheels: macOS (Apple Silicon), Linux (x86_64,
16
+ aarch64) and Windows (x64), Python 3.10+.
17
+
18
+ ## Usage
19
+
20
+ ```python
21
+ from offline_translate import Translator
22
+
23
+ t = Translator("~/.local/share/offline-translate") # where models are stored
24
+
25
+ t.download("he") # Hebrew <-> English, ~55 MB
26
+ t.translate("Good morning!", target="he").text
27
+ # 'בוקר טוב!'
28
+
29
+ t.translate("בוקר טוב", target="en") # source is auto-detected among installed languages
30
+ # Translation(text='Good morning', source='iw', target='en', via=('iw', 'en'))
31
+ ```
32
+
33
+ ### Managing languages
34
+
35
+ ```python
36
+ for lang in t.languages(): # full catalog (fetched once from Google, then cached)
37
+ print(lang.code, lang.name, lang.installed, lang.download_size, lang.update_available)
38
+
39
+ t.installed_languages() # ['en', 'iw'] — works offline, no catalog needed
40
+ t.download("fr", progress=lambda done, total: print(f"{done}/{total}"))
41
+ t.remove("fr")
42
+
43
+ t.update_catalog() # check Google for a newer catalog
44
+ [l for l in t.languages() if l.update_available] # then t.download(code) to update
45
+ ```
46
+
47
+ Language codes follow Google's catalog (`iw` for Hebrew, `zh-Hant`, `jw`, …). Common aliases
48
+ (`he`, `zh-TW`, `jv`) and English names (`"Hebrew"`) are accepted everywhere.
49
+
50
+ English is built in: there is no separate English pack, since every pack contains both
51
+ directions to and from English.
52
+
53
+ ### API
54
+
55
+ | | |
56
+ |---|---|
57
+ | `Translator(models_dir, channel="prod", beam=4)` | `models_dir` is created if missing. `beam=1` is faster, slightly worse. |
58
+ | `translate(text, target, source=None) -> Translation` | `Translation(text, source, target, via)` |
59
+ | `detect(text) -> str` | Guess among installed languages |
60
+ | `languages(installed_only=False) -> list[Language]` | `Language(code, name, installed, download_size, size, update_available, builtin, supported)` |
61
+ | `installed_languages() -> list[str]` | |
62
+ | `download(lang, progress=None, force=False)` | Verifies SHA-1 of every file; resumes interrupted downloads |
63
+ | `remove(lang)` | |
64
+ | `update_catalog() -> str` | Returns the catalog version |
65
+ | `TranslateError` | Raised for unknown/uninstalled languages, checksum failures, … |
66
+
67
+ Network errors surface as `OSError`.
68
+
69
+ ## Command line
70
+
71
+ The repository also has a small CLI built on the library (not part of the PyPI package):
72
+
73
+ ```bash
74
+ ./translate get he # download Hebrew
75
+ ./translate "Good morning" # English -> Hebrew by default; source auto-detected
76
+ ./translate -s he -t fr "בוקר טוב" # any pair, via English
77
+ ./translate list --installed
78
+ ```
79
+
80
+ `tools/pull_from_phone.py` copies already-downloaded packs from a rooted Android phone over adb.
81
+
82
+ ## How it works
83
+
84
+ The Android app reads a manifest at
85
+ `https://dl.google.com/translate/offline/manifest/gtr/prod/manifest.pb`, which lists versions of
86
+ a *package profile* — the catalog of language packs and their files. Each translation direction is
87
+ a TFLite encoder, an LSTM decoder run step by step with beam search, and SentencePiece tokenizers.
88
+
89
+ ### Limitations
90
+
91
+ - 15 languages (am, az, bs, hy, kk, ky, ml, mn, ne, pa, ps, si, so, tg, uz) still ship an older
92
+ model format with a WordPiece tokenizer that isn't supported yet; they're listed with
93
+ `supported=False`.
94
+ - Quality matches the app's offline mode, which is below online Google Translate.
95
+
96
+ ## Development
97
+
98
+ ```bash
99
+ pip install -e ".[test]"
100
+ pytest # offline unit tests (fake Google server, fake model)
101
+ pytest --network # + checks against the real Google catalog
102
+ pytest --models=~/models # + end-to-end translation with installed Hebrew models
103
+ ```
104
+
105
+ ## Disclaimer
106
+
107
+ This project is not affiliated with or endorsed by Google. It contains no Google code or models;
108
+ language packs are downloaded from Google's public servers at runtime and are subject to Google's
109
+ terms.
@@ -0,0 +1,8 @@
1
+ """Root conftest: pytest only reads command-line options registered here."""
2
+ import os
3
+
4
+
5
+ def pytest_addoption(parser):
6
+ parser.addoption("--network", action="store_true", help="run tests that talk to dl.google.com")
7
+ parser.addoption("--models", metavar="DIR", default=os.environ.get("OFFLINE_TRANSLATE_TEST_MODELS"),
8
+ help="models folder with Hebrew installed, for end-to-end translation tests")
@@ -0,0 +1,55 @@
1
+ [build-system]
2
+ requires = ["hatchling>=1.27"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "offline-translate"
7
+ dynamic = ["version"]
8
+ description = "Offline machine translation using Google Translate's on-device language packs"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ authors = [{ name = "asimon655" }]
14
+ keywords = ["translation", "offline", "nmt", "tflite", "google-translate"]
15
+ classifiers = [
16
+ "Development Status :: 3 - Alpha",
17
+ "Intended Audience :: Developers",
18
+ "Operating System :: MacOS",
19
+ "Operating System :: POSIX :: Linux",
20
+ "Operating System :: Microsoft :: Windows",
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3 :: Only",
23
+ "Topic :: Text Processing :: Linguistic",
24
+ "Typing :: Typed",
25
+ ]
26
+ dependencies = [
27
+ "ai-edge-litert>=1.0",
28
+ "numpy>=1.24",
29
+ "sentencepiece>=0.2",
30
+ ]
31
+
32
+ [project.urls]
33
+ Homepage = "https://github.com/asimon655/offline-translate"
34
+ Issues = "https://github.com/asimon655/offline-translate/issues"
35
+
36
+ [project.optional-dependencies]
37
+ test = ["pytest>=8"]
38
+
39
+ [tool.hatch.version]
40
+ path = "src/offline_translate/__init__.py"
41
+
42
+ [tool.hatch.build.targets.wheel]
43
+ packages = ["src/offline_translate"]
44
+
45
+ # Explicit allow-list: models/, phone_dump/ etc. must never end up in an upload.
46
+ [tool.hatch.build.targets.sdist]
47
+ only-include = ["src", "tests", "conftest.py", "README.md", "LICENSE", "pyproject.toml"]
48
+
49
+ [tool.pytest.ini_options]
50
+ testpaths = ["tests"]
51
+ addopts = "-ra"
52
+ markers = [
53
+ "network: talks to dl.google.com (enable with --network)",
54
+ "models: needs real installed models (enable with --models=DIR)",
55
+ ]
@@ -0,0 +1,12 @@
1
+ """Offline translation with Google Translate's on-device language packs.
2
+
3
+ from offline_translate import Translator
4
+ t = Translator("~/models")
5
+ t.download("he")
6
+ t.translate("Hello", target="he").text
7
+ """
8
+ from .languages import NAMES, normalize
9
+ from .translator import Language, TranslateError, Translation, Translator
10
+
11
+ __version__ = "0.1.0"
12
+ __all__ = ["Translator", "Translation", "Language", "TranslateError", "NAMES", "normalize", "__version__"]
@@ -0,0 +1,172 @@
1
+ """Google Translate offline package catalog.
2
+
3
+ How the Android app finds packs (reverse engineered):
4
+ 1. Manifest https://dl.google.com/translate/offline/manifest/gtr/<channel>/manifest.pb
5
+ lists every package-profile version (+ sha1/size). The channel comes from a server flag ("prod").
6
+ 2. Profile https://dl.google.com/translate/offline/package_profile/<maj>/<min>/<patch>/translate_premium.pb.zlib
7
+ lists every language pack, its packages (nmt models, dictionary, translit) and their files.
8
+ 3. Files https://dl.google.com/translate/offline/nmt/<ver>/premium/<src>_<tgt>/0/<file>.zlib
9
+ """
10
+ import struct
11
+ import urllib.request
12
+ import zlib
13
+ from dataclasses import dataclass, field
14
+
15
+ BASE = "https://dl.google.com/translate/offline"
16
+ MANIFEST_URL = BASE + "/manifest/gtr/{channel}/manifest.pb"
17
+ USER_AGENT = "Mozilla/5.0"
18
+
19
+ PKG_NMT = 3
20
+ # Files our inference engine knows how to run; packages with other layouts are skipped.
21
+ NMT_FILES = {
22
+ "encoder_0.tflite", "decoder_init_0.tflite", "decoder_step_0.tflite",
23
+ "src_spm_model.model", "tgt_spm_model.model",
24
+ }
25
+
26
+
27
+ # --- minimal protobuf reader (Google's schema isn't public) ---
28
+ def _varint(b, i):
29
+ r = s = 0
30
+ while True:
31
+ c = b[i]
32
+ i += 1
33
+ r |= (c & 0x7F) << s
34
+ s += 7
35
+ if c < 0x80:
36
+ return r, i
37
+
38
+
39
+ def parse_pb(b):
40
+ """Parse a protobuf message into {field: [values]}; length-delimited values stay bytes."""
41
+ out, i = {}, 0
42
+ while i < len(b):
43
+ key, i = _varint(b, i)
44
+ f, t = key >> 3, key & 7
45
+ if t == 0:
46
+ v, i = _varint(b, i)
47
+ elif t == 2:
48
+ n, i = _varint(b, i)
49
+ v, i = b[i:i + n], i + n
50
+ elif t == 5:
51
+ v, i = struct.unpack("<f", b[i:i + 4])[0], i + 4
52
+ elif t == 1:
53
+ v, i = struct.unpack("<d", b[i:i + 8])[0], i + 8
54
+ else:
55
+ raise ValueError(f"unsupported protobuf wire type {t}")
56
+ out.setdefault(f, []).append(v)
57
+ return out
58
+
59
+
60
+ def _str(m, f):
61
+ return m[f][0].decode()
62
+
63
+
64
+ def _hash(m):
65
+ """{1: {1: algo, 2: hex}, 2: size} -> (hex, size)."""
66
+ return _str(parse_pb(m[1][0]), 2), m[2][0]
67
+
68
+
69
+ def fetch(url, timeout=60):
70
+ req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
71
+ with urllib.request.urlopen(req, timeout=timeout) as r:
72
+ return r.read()
73
+
74
+
75
+ def open_url(url, timeout=60):
76
+ return urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT}), timeout=timeout)
77
+
78
+
79
+ @dataclass(frozen=True)
80
+ class ProfileRef:
81
+ version: tuple
82
+ url: str
83
+ sha1: str # of the decompressed profile
84
+
85
+ @property
86
+ def version_str(self):
87
+ return ".".join(map(str, self.version))
88
+
89
+
90
+ def parse_manifest(data):
91
+ """Returns the premium (NMT) package profiles listed in a manifest, oldest first."""
92
+ refs = []
93
+ groups = parse_pb(data).get(2, [])
94
+ for raw in (r for g in groups for r in parse_pb(g).get(2, [])):
95
+ m = parse_pb(raw)
96
+ url = _str(m, 2)
97
+ if "translate_premium" not in url:
98
+ continue
99
+ v = parse_pb(m[1][0])
100
+ version = tuple(v.get(k, [0])[0] for k in (1, 2, 3))
101
+ refs.append(ProfileRef(version, url, _hash(parse_pb(m[3][0]))[0]))
102
+ return sorted(refs, key=lambda r: r.version)
103
+
104
+
105
+ @dataclass(frozen=True)
106
+ class PackFile:
107
+ name: str # file name on disk, e.g. encoder_0.tflite
108
+ url: str
109
+ sha1: str # of the decompressed file
110
+ size: int # decompressed bytes
111
+ download_size: int # zlib-compressed bytes
112
+
113
+
114
+ @dataclass(frozen=True)
115
+ class Package:
116
+ name: str # e.g. "en_iw" (direction src_tgt)
117
+ files: tuple
118
+
119
+ @property
120
+ def source(self):
121
+ return self.name.split("_", 1)[0]
122
+
123
+ @property
124
+ def target(self):
125
+ return self.name.split("_", 1)[1]
126
+
127
+
128
+ @dataclass
129
+ class Pack:
130
+ """A downloadable language: both directions to/from English."""
131
+ lang: str
132
+ packages: list = field(default_factory=list)
133
+ # False for older packs (nmt/8/1, WordPiece tokenizer) that the engine can't run yet.
134
+ supported: bool = True
135
+
136
+ @property
137
+ def download_size(self):
138
+ return sum(f.download_size for p in self.packages for f in p.files)
139
+
140
+ @property
141
+ def size(self):
142
+ return sum(f.size for p in self.packages for f in p.files)
143
+
144
+
145
+ def parse_profile(data):
146
+ """Returns {lang_code: Pack} for every language with an NMT model (see Pack.supported)."""
147
+ top = parse_pb(data)
148
+ files = {}
149
+ for raw in top.get(3, []):
150
+ m = parse_pb(raw)
151
+ sha1, size = _hash(parse_pb(m[3][0]))
152
+ _, zsize = _hash(parse_pb(m[4][0]))
153
+ files[_str(m, 1)] = PackFile(_str(m, 5), _str(m, 2), sha1, size, zsize)
154
+
155
+ packages = {}
156
+ for raw in top.get(2, []):
157
+ m = parse_pb(raw)
158
+ if m[2][0] != PKG_NMT:
159
+ continue
160
+ pfiles = tuple(files[f.decode()] for f in m.get(4, []))
161
+ packages[_str(m, 1)] = Package(_str(m, 1).removesuffix("_premium_nmt"), pfiles)
162
+
163
+ packs = {}
164
+ for raw in top.get(1, []):
165
+ m = parse_pb(raw)
166
+ langs = [_str(parse_pb(x), 2) for x in parse_pb(m[3][0])[1]]
167
+ lang = next((l for l in langs if l != "en"), None)
168
+ pkgs = [packages[p.decode()] for p in m.get(2, []) if p.decode() in packages]
169
+ if lang and pkgs:
170
+ ok = all(NMT_FILES <= {f.name for f in p.files} for p in pkgs)
171
+ packs[lang] = Pack(lang, pkgs, ok)
172
+ return packs
@@ -0,0 +1,118 @@
1
+ """Runs one direction of a Google Translate on-device NMT model (TFLite encoder + LSTM decoder)."""
2
+ import os
3
+ import re
4
+ from pathlib import Path
5
+
6
+ os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
7
+
8
+ import numpy as np
9
+ import sentencepiece as spm
10
+ from ai_edge_litert.interpreter import Interpreter
11
+
12
+ BOS, EOS = 1, 2
13
+ STATE_KEYS = [f"beam_state:rnn_state_{i}_{x}:00" for i in (0, 1) for x in ("c", "m")]
14
+ # The models are trained on single sentences; split on sentence-final punctuation + whitespace.
15
+ SENT_RE = re.compile(r"(?<=[.!?…。!?])\s+")
16
+
17
+
18
+ def split_sentences(text):
19
+ """Lines of sentences: line breaks are kept, each line is split into sentences."""
20
+ return [[s for s in SENT_RE.split(line.strip()) if s] for line in text.split("\n")]
21
+
22
+
23
+ FINAL_PUNCT = ".!?…。!?:;\"'”’)»"
24
+
25
+
26
+ def add_period_limit(processor_spec):
27
+ """max_sent_len of the pack's PreprocAddPeriod step, or None if the pack doesn't use it."""
28
+ m = re.search(rb"PreprocAddPeriod/[^:]*:max_sent_len=(\d+)", processor_spec)
29
+ return int(m.group(1)) if m else None
30
+
31
+
32
+ def with_period(sentence, limit):
33
+ """Short sentences without final punctuation translate badly ("car" -> "tram"); the app adds
34
+ a period first. Returns (text to translate, whether a period was added)."""
35
+ if limit and len(sentence.split()) <= limit and sentence[-1:] not in FINAL_PUNCT:
36
+ return sentence + ".", True
37
+ return sentence, False
38
+
39
+
40
+ def strip_period(translation):
41
+ return translation[:-1] if translation[-1:] in ".。" else translation
42
+
43
+
44
+ class NmtModel:
45
+ def __init__(self, path):
46
+ path = Path(path)
47
+ self.src_sp = spm.SentencePieceProcessor(model_file=str(path / "src_spm_model.model"))
48
+ self.tgt_sp = spm.SentencePieceProcessor(model_file=str(path / "tgt_spm_model.model"))
49
+ runner = lambda name: Interpreter(model_path=str(path / f"{name}.tflite")).get_signature_runner()
50
+ self.encoder = runner("encoder_0")
51
+ self.dec_init = runner("decoder_init_0")
52
+ self.dec_step = runner("decoder_step_0")
53
+ spec = path / "processor_spec.pb"
54
+ self.period_limit = add_period_limit(spec.read_bytes()) if spec.exists() else None
55
+
56
+ def translate(self, text, beam=4):
57
+ return "\n".join(" ".join(self.translate_sentence(s, beam=beam) for s in line)
58
+ for line in split_sentences(text))
59
+
60
+ def translate_sentence(self, text, beam=4):
61
+ text, added = with_period(text, self.period_limit)
62
+ out = self._decode(text, beam=beam)
63
+ return strip_period(out) if added else out
64
+
65
+ def _decode(self, text, beam=4, max_len=72, alpha=1.0):
66
+ # Source vocabularies are lowercase-only; the model restores casing on output.
67
+ ids = self.src_sp.encode(text.lower()) + [EOS]
68
+ src = np.array([ids], np.int32)
69
+ pad = np.zeros_like(src, np.float32)
70
+ encs = self.encoder(**{"feeds:src_ids:00": src, "feeds:src_paddings:00": pad})["fetches:src_encs:00"]
71
+ init = self.dec_init(**{
72
+ "feeds:num_hyps:00": np.array([1], np.int32),
73
+ "feeds:src_enc_paddings:00": pad,
74
+ "feeds:src_encs:00": encs,
75
+ })
76
+ vecs = init["atten_init:source_vecs:00"]
77
+ state = {k: np.zeros((1, 512), np.float32) for k in STATE_KEYS}
78
+ state["beam_state:atten_context:00"] = init["beam_state:atten_context:00"]
79
+
80
+ # Each live hypothesis: (tokens, cumulative log prob). Batch rows of `state` align with `hyps`.
81
+ hyps = [([], 0.0)]
82
+ done = []
83
+ max_len = min(max_len, 2 * len(ids) + 10)
84
+ for _ in range(max_len):
85
+ n = len(hyps)
86
+ last = np.array([[t[-1] if t else BOS] for t, _ in hyps], np.int32)
87
+ out = self.dec_step(**{
88
+ "atten_init:source_contexts:00": np.repeat(vecs, n, 0),
89
+ "atten_init:source_vecs:00": np.repeat(vecs, n, 0),
90
+ "atten_init:source_padding:00": np.repeat(pad, n, 0),
91
+ "feeds:step_ids:00": last,
92
+ **state,
93
+ })
94
+ logp = out["fetches:log_probs:00"]
95
+ scores = np.array([s for _, s in hyps])[:, None] + logp
96
+ flat = np.argsort(scores, axis=None)[::-1][: 2 * beam]
97
+ new_hyps, rows = [], []
98
+ for idx in flat:
99
+ row, tok = divmod(int(idx), logp.shape[1])
100
+ toks, sc = hyps[row][0] + [tok], float(scores[row, tok])
101
+ if tok == EOS:
102
+ done.append((toks[:-1], sc / (((5 + len(toks)) / 6) ** alpha)))
103
+ elif len(new_hyps) < beam:
104
+ new_hyps.append((toks, sc))
105
+ rows.append(row)
106
+ if not new_hyps:
107
+ break
108
+ # Stop when the best finished hypothesis beats any live one's optimistic score.
109
+ if len(done) >= beam:
110
+ best_done = max(s for _, s in done)
111
+ best_live = max(s for _, s in new_hyps) / (((5 + max_len) / 6) ** alpha)
112
+ if best_done > best_live:
113
+ break
114
+ hyps = new_hyps
115
+ state = {k: out[k][rows] for k in state}
116
+ if not done:
117
+ done = hyps
118
+ return self.tgt_sp.decode(max(done, key=lambda h: h[1])[0])