| """ |
| Core inference engine for Hugging Face Space. |
| Downloads bundles on-demand from the Hugging Face Hub (or local directory) |
| and maintains an in-memory LRU cache of loaded models. |
| """ |
|
|
| import json |
| import logging |
| import os |
| from functools import lru_cache |
| from typing import Dict, List, Optional, Tuple |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
|
|
| from demo_core.loader import load_bundle |
| from demo_core.predict import PredictionResult, predict_single |
|
|
| HF_MODEL_REPO = os.environ.get("HF_MODEL_REPO", "simon-clmtd/neural-transducer-inflection") |
| LOCAL_MODELS_DIR = os.environ.get("LOCAL_MODELS_DIR", "models") |
|
|
|
|
| class TransducerSpaceEngine: |
| def __init__(self, model_repo: str = HF_MODEL_REPO, local_dir: str = LOCAL_MODELS_DIR): |
| self.model_repo = model_repo |
| self.local_dir = local_dir |
| self._cache: Dict[str, Tuple] = {} |
| self._cache_order: List[str] = [] |
| self.max_cache_size = 3 |
| self.device = "cpu" |
|
|
| def _resolve_bundle_path(self, bundle_name: str) -> str: |
| |
| local_path = os.path.join(self.local_dir, bundle_name) |
| if os.path.exists(os.path.join(local_path, "best.model.json")): |
| return local_path |
|
|
| |
| logging.info("Downloading bundle %s from %s...", bundle_name, self.model_repo) |
| download_dir = snapshot_download( |
| repo_id=self.model_repo, |
| allow_patterns=f"{bundle_name}/*", |
| local_dir_use_symlinks=False, |
| ) |
| target_dir = os.path.join(download_dir, bundle_name) |
| if not os.path.exists(target_dir): |
| raise FileNotFoundError(f"Bundle {bundle_name} not found after download in {download_dir}") |
| return target_dir |
|
|
| def get_bundle(self, lang: str, regime: str) -> Tuple: |
| bundle_name = f"{lang}_{regime}" |
| if bundle_name in self._cache: |
| |
| self._cache_order.remove(bundle_name) |
| self._cache_order.append(bundle_name) |
| return self._cache[bundle_name] |
|
|
| |
| if len(self._cache_order) >= self.max_cache_size: |
| oldest = self._cache_order.pop(0) |
| del self._cache[oldest] |
| logging.info("Evicted bundle %s from LRU cache.", oldest) |
|
|
| path = self._resolve_bundle_path(bundle_name) |
| model, vocab, meta = load_bundle(path, device=self.device) |
| self._cache[bundle_name] = (model, vocab, meta) |
| self._cache_order.append(bundle_name) |
| return model, vocab, meta |
|
|
| def predict( |
| self, |
| lang: str, |
| regime: str, |
| lemma: str, |
| features: str, |
| beam_width: int = 1, |
| ) -> PredictionResult: |
| model, vocab, _ = self.get_bundle(lang, regime) |
| return predict_single( |
| model=model, |
| vocab=vocab, |
| lemma=lemma, |
| features_str=features, |
| beam_width=beam_width, |
| ) |
|
|
| def _load_data_index(self) -> Dict: |
| if not hasattr(self, "_data_index"): |
| index_path = os.path.join(os.path.dirname(__file__), "data_index.json") |
| if not os.path.exists(index_path): |
| |
| index_path = os.path.join(os.path.dirname(__file__), "..", "demo_core", "data_index.json") |
| if os.path.exists(index_path): |
| with open(index_path, "r", encoding="utf-8") as f: |
| self._data_index = json.load(f) |
| else: |
| self._data_index = {} |
| return self._data_index |
|
|
| def get_random_sample( |
| self, lang: str, split: str = "dev", filter_type: str = "any" |
| ) -> Optional[Dict[str, str]]: |
| import random |
| data = self._load_data_index() |
| split_key = "dev" if split.lower().startswith("dev") else "test" |
| samples = data.get(lang, {}).get(split_key, []) |
| if not samples: |
| return None |
|
|
| if filter_type == "same": |
| filtered = [s for s in samples if s.get("lemma") == s.get("target")] |
| elif filter_type == "different": |
| filtered = [s for s in samples if s.get("lemma") != s.get("target")] |
| else: |
| filtered = samples |
|
|
| if not filtered: |
| return None |
| return random.choice(filtered) |
|
|
| def check_training_coverage(self, lang: str, lemma: str, features_str: str) -> Dict[str, Dict[str, bool]]: |
| data = self._load_data_index() |
| lang_data = data.get(lang, {}).get("training_coverage", {}) |
| lemma = lemma.strip() |
| features_norm = set(features_str.strip().split(";")) |
|
|
| result = {} |
| for regime in ["100", "1000", "10000"]: |
| cov = lang_data.get(regime, {}) |
| seen_lemmas = set(cov.get("lemmas", [])) |
| raw_feats = cov.get("features", []) |
| seen_feats_sets = [set(f.split(";")) for f in raw_feats] |
|
|
| lemma_seen = lemma in seen_lemmas |
| |
| feats_seen = (features_str.strip() in raw_feats) or (features_norm in seen_feats_sets) |
|
|
| result[regime] = { |
| "lemma_seen": lemma_seen, |
| "features_seen": feats_seen, |
| } |
| return result |
|
|
|
|
