82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
#!/usr/bin/env python
|
|
"""Fetch the built-in embedding model and check it against the manifest.
|
|
|
|
Runs at build time, never on a user's machine: the point of building the
|
|
model in is that a first launch works with no network. Downloading here and
|
|
verifying against `embedding_assets.EMBEDDING_ASSETS` is what lets the weights
|
|
stay out of a repository that is about to go public.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT / "apps" / "core" / "src"))
|
|
|
|
from one_memory.embedding_assets import ( # noqa: E402
|
|
EMBEDDING_ASSETS,
|
|
AssetVerificationError,
|
|
digest_of,
|
|
resolve_asset_dir,
|
|
verify_assets,
|
|
)
|
|
|
|
|
|
def fetch(target: Path, *, force: bool) -> int:
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
for asset in EMBEDDING_ASSETS:
|
|
path = target / asset.filename
|
|
if path.is_file() and not force and digest_of(path) == asset.sha256:
|
|
print(f" ok {asset.filename}")
|
|
continue
|
|
print(f" fetching {asset.filename} from {asset.url}")
|
|
with urllib.request.urlopen(asset.url, timeout=300) as response:
|
|
path.write_bytes(response.read())
|
|
|
|
try:
|
|
verify_assets(target)
|
|
except AssetVerificationError as error:
|
|
# Do not leave unverified bytes where the encoder will happily load
|
|
# them: a partial download that still parses is the failure mode this
|
|
# whole manifest exists to catch.
|
|
print(f"\nverification failed: {error}", file=sys.stderr)
|
|
for asset in EMBEDDING_ASSETS:
|
|
(target / asset.filename).unlink(missing_ok=True)
|
|
return 1
|
|
|
|
total = sum((target / asset.filename).stat().st_size for asset in EMBEDDING_ASSETS)
|
|
print(f"\nverified {len(EMBEDDING_ASSETS)} files, {total / 1048576:.1f} MB")
|
|
print(f"into {target}")
|
|
return 0
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--target", type=Path, default=None)
|
|
parser.add_argument("--force", action="store_true")
|
|
parser.add_argument(
|
|
"--verify-only",
|
|
action="store_true",
|
|
help="check what is already on disk and fetch nothing",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
target = args.target or resolve_asset_dir()
|
|
if args.verify_only:
|
|
try:
|
|
verify_assets(target)
|
|
except AssetVerificationError as error:
|
|
print(f"verification failed: {error}", file=sys.stderr)
|
|
return 1
|
|
print(f"verified {len(EMBEDDING_ASSETS)} files in {target}")
|
|
return 0
|
|
return fetch(target, force=args.force)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|