File size: 5,441 Bytes
ff372d4 3746f17 ff372d4 0f9f04c 067607f 0f9f04c 067607f 0f9f04c 3746f17 0f9f04c | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
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" # Space runs on cpu-basic
def _resolve_bundle_path(self, bundle_name: str) -> str:
# Check local path first (useful during local testing and development)
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
# Download from Hugging Face Hub snapshot
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:
# Move to end of LRU
self._cache_order.remove(bundle_name)
self._cache_order.append(bundle_name)
return self._cache[bundle_name]
# Evict oldest if cache limit reached
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):
# Fallback to demo_core location if needed
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
# Feat combination seen if exact string matches or tag set matches
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
|