| """ |
| Metadata-driven model loader for neural transducer bundles. |
| Reconstructs Transducer architecture exclusively from best.model.json metadata, |
| ensuring zero configuration hard-coding in client apps or scripts. |
| """ |
|
|
| import argparse |
| import hashlib |
| import json |
| import logging |
| import os |
| from typing import Optional, Tuple |
|
|
| import torch |
|
|
| from trans import optimal_expert_substitutions |
| from trans import sed |
| from trans import transducer |
| from trans import vocabulary |
|
|
|
|
| def compute_sha256(filepath: str) -> str: |
| sha = hashlib.sha256() |
| with open(filepath, "rb") as f: |
| for chunk in iter(lambda: f.read(65536), b""): |
| sha.update(chunk) |
| return sha.hexdigest() |
|
|
|
|
| def verify_bundle_manifest(bundle_dir: str) -> dict: |
| """Verifies that all files in bundle_dir match bundle_manifest.json hashes.""" |
| manifest_path = os.path.join(bundle_dir, "bundle_manifest.json") |
| if not os.path.exists(manifest_path): |
| raise FileNotFoundError(f"Missing bundle_manifest.json in {bundle_dir}") |
|
|
| with open(manifest_path, "r", encoding="utf-8") as f: |
| manifest = json.load(f) |
|
|
| for filename, expected_hash in manifest.get("files", {}).items(): |
| filepath = os.path.join(bundle_dir, filename) |
| if not os.path.exists(filepath): |
| raise FileNotFoundError(f"Bundle file missing: {filepath}") |
| actual_hash = compute_sha256(filepath) |
| if actual_hash != expected_hash: |
| raise ValueError( |
| f"Bundle file integrity compromised for {filename}: " |
| f"expected {expected_hash}, got {actual_hash}" |
| ) |
| return manifest |
|
|
|
|
| def load_bundle(bundle_dir: str, device: Optional[str] = None) -> Tuple[transducer.Transducer, vocabulary.FeatureVocabularies, dict]: |
| """ |
| Loads a fine-tuned Transducer model bundle. |
| |
| Returns: |
| (model, vocabulary, metadata) |
| """ |
| if device is None: |
| device = "mps" if torch.backends.mps.is_available() else "cpu" |
|
|
| metadata_path = os.path.join(bundle_dir, "best.model.json") |
| model_path = os.path.join(bundle_dir, "best.model") |
| voc_path = os.path.join(bundle_dir, "vocabulary.pkl") |
| sed_path = os.path.join(bundle_dir, "sed.pkl") |
|
|
| for path in (metadata_path, model_path, voc_path, sed_path): |
| if not os.path.exists(path): |
| raise FileNotFoundError(f"Required bundle artifact not found: {path}") |
|
|
| |
| manifest_path = os.path.join(bundle_dir, "bundle_manifest.json") |
| if os.path.exists(manifest_path): |
| verify_bundle_manifest(bundle_dir) |
|
|
| |
| with open(metadata_path, "r", encoding="utf-8") as f: |
| metadata = json.load(f) |
|
|
| train_args_dict = metadata.get("args", {}) |
| |
| train_args_dict["device"] = device |
| args = argparse.Namespace(**train_args_dict) |
|
|
| |
| vocab = vocabulary.FeatureVocabularies.from_pickle(voc_path) |
| sed_aligner = sed.StochasticEditDistance.from_pickle(sed_path) |
| expert = optimal_expert_substitutions.OptimalSubstitutionExpert(sed_aligner) |
|
|
| |
| model = transducer.Transducer(vocab, expert, args) |
|
|
| |
| state_dict = torch.load(model_path, map_location=device) |
| model.load_state_dict(state_dict) |
| model.to(device) |
| model.eval() |
|
|
| logging.info("Successfully loaded bundle from %s to device %s", bundle_dir, device) |
| return model, vocab, metadata |
|
|
