| """Arc Tagger: text-only streaming tagger (about 10 M parameters, CPU) for ASR word streams. |
| |
| from arctag import ArcTagger |
| tg = ArcTagger.from_pretrained() # WhissleAI/arc-tagger from the Hub, or a local dir |
| out = tg.tag("Hi, this is Dana from Acme. Can you move my appointment to Friday?") |
| |
| st = tg.stream(domain="talk") # streaming: words in as the ASR emits them |
| st.push_words(["hi", "this", "is"], final=False) # interim |
| st.push_words(["hi,", "this", "is", "Dana."], revise_from=0, final=True) |
| st.endpoint() |
| res = st.finalize() |
| |
| CPU only. The model code is the streaming engine in arctag/_core. |
| """ |
| import importlib.util, json, os, sys |
|
|
| __version__ = "0.4.1" |
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| _CORE = os.path.join(_HERE, "_core") |
| _EN = None |
|
|
|
|
| def _engine(home): |
| """import the vendored engine once (module-level path constants resolve against `home`)""" |
| global _EN |
| if _EN is None: |
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") |
| os.environ.setdefault("ARCTAG_HOME", home) |
| os.environ["ARCTAG_ROLE"] = "offline" |
| if _CORE not in sys.path: |
| sys.path.insert(0, _CORE) |
| spec = importlib.util.spec_from_file_location("arctag_engine", os.path.join(_CORE, "engine.py")) |
| _EN = importlib.util.module_from_spec(spec); sys.modules["arctag_engine"] = _EN |
| spec.loader.exec_module(_EN) |
| return _EN |
|
|
|
|
| def _driver(): |
| from . import driver |
| return driver |
|
|
|
|
| class ArcTagger: |
| def __init__(self, home, EN, loaded, cfg): |
| self.home, self.EN, self.loaded, self.config = home, EN, loaded, cfg |
|
|
| @classmethod |
| def from_pretrained(cls, path="WhissleAI/arc-tagger", revision=None, thresholds="refit"): |
| """path: release dir (model.safetensors, config.json, tok/, ontology/, thresholds/) or a Hugging Face repo id. |
| thresholds: "refit" (default: thresholds/display.json, plus thresholds/events.json and span_head.json when present) |
| or "none" (the default thresholds in config.json only; engine display floors still apply).""" |
| if not os.path.isdir(path): |
| from huggingface_hub import snapshot_download |
| path = snapshot_download(path, revision=revision, allow_patterns=["*.json", "*.safetensors", "*.npy", "*.py", "*.txt", "*.md", |
| "LICENSE*", "tok/*", "ontology/*", "thresholds/*", "arctag/*"]) |
| home = os.path.abspath(path) |
| EN = _engine(home) |
| AT = sys.modules["arc_text"]; TT = sys.modules["tagtable"] |
| AT.BPE8K = os.path.join(home, "tok", "bpe8k", "tokenizer.json") |
| cfg = json.load(open(os.path.join(home, "config.json"))) |
| ext = os.path.join(home, cfg.get("onto_ext") or "ontology/ext_v7_aff.json") |
| T = TT.TextTable(r8=os.path.join(home, "ontology", "tags_v1_r8.json"), |
| text_tags=os.path.join(home, "ontology", "text_tags.json"), ext=ext) |
| T.load_E(os.path.join(home, "ontology", "head_E_r8_bge_small.npy")) |
| if T.sha != cfg.get("ontology_sha"): |
| raise RuntimeError(f"ontology sha {T.sha} != config {cfg.get('ontology_sha')}") |
| loaded = _load_model(EN, home, cfg, T, thresholds) |
| return cls(home, EN, loaded, cfg) |
|
|
| def stream(self, domain="auto", tier="T0"): |
| """a streaming session (driver.Stream): push_words(words, final=True, revise_from=None, endpoint=False), endpoint(), |
| finalize() -> result dict""" |
| return _driver().Stream(self.EN, self.loaded, domain=domain, tier=tier, meta=self._meta()) |
|
|
| def tag(self, x, domain="auto", tier="T0"): |
| """x: text (str), a list of words (str) or word dicts {"w", "t0", "t1", "conf", "spk", "final"}. Utterances end after |
| words ending in . ? ! Returns the structured result (see README "Output").""" |
| return _driver().tag(self.EN, self.loaded, x, domain=domain, tier=tier, meta=self._meta()) |
|
|
| def _meta(self): |
| return {"model": "arc-tagger", "version": __version__} |
|
|
|
|
| def _read_thr(display, events): |
| """thresholds/display.json {"thresholds": {tag@readout: {domain|_all: thr}}} and thresholds/events.json {"tier", "hide", |
| "thresholds"}""" |
| pref, pstep, ev, evstep = {}, None, {}, None |
| if os.path.exists(display): |
| pref, pstep = json.load(open(display)).get("thresholds") or {}, "display" |
| if os.path.exists(events): |
| ev, evstep = json.load(open(events)), "events" |
| return pref, pstep, ev, evstep |
|
|
|
|
| def _load_model(EN, home, cfg, T, thresholds="refit"): |
| """load model.safetensors and the threshold tables of a release dir""" |
| from safetensors.torch import load_file |
| AT = sys.modules["arc_text"]; SA = sys.modules["stream_arc"] |
| step = None |
| m = AT.ArcText(T, scratch=cfg.get("scratch") or AT.SCRATCH_10M) |
| sd = load_file(os.path.join(home, "model.safetensors")) |
| SA.add_modules_for(m, sd) |
| missing, unexpected = m.load_state_dict(sd, strict=False) |
| if unexpected: |
| raise RuntimeError(f"weights without a module: {unexpected[:6]}") |
| bad = [k for k in missing if not k.startswith(("head.E",))] |
| if bad: |
| import warnings |
| warnings.warn(f"{len(bad)} missing keys (e.g. {bad[:3]})") |
| m.eval() |
| own = cfg.get("thresholds") or {} |
| pref, pstep, ev, evstep = {}, None, {}, None |
| if thresholds == "refit": |
| d = os.path.join(home, "thresholds") |
| pref, pstep, ev, evstep = _read_thr(os.path.join(d, "display.json"), os.path.join(d, "events.json")) |
| sh = os.path.join(d, "span_head.json") |
| if os.path.exists(sh): |
| ts = {t: dict(v) for t, v in (cfg.get("thresholds_span") or {}).items()} |
| for t, v in (json.load(open(sh)).get("thresholds_span") or {}).items(): |
| ts.setdefault(t, {}).update(v) |
| cfg = dict(cfg, thresholds_span=ts) |
| elif thresholds != "none": |
| raise ValueError("thresholds must be 'refit' or 'none'") |
| thr = EN.Thresholds(own, {}, None, pref=pref, pref_step=pstep) |
| thr.set_events(ev, evstep) |
| L = EN.Loaded(step, home, m, T, thr, cfg) |
| fams = cfg.get("span_head_families") |
| if fams is not None and L.head_tags: |
| L.head_use, L.head_why = set(fams) & set(L.head_tags), "config" |
| return L |
|
|
