#!/usr/bin/env python3 """ sphinx_trie.py — Gardiner-keyed prefix trie for the SphinxEyes inference pipeline. Role: validate / correct YOLOv11 hieroglyph predictions before passing the phonetic transliteration to GPT for translation. Edges = single Gardiner codes (your 148 trained classes + 'Unknown'). Terminal nodes carry a list of SourceRecord, one per observation in the source corpora (Dickson dictionary, BBAW corpus, TLA corpus). Module-level class so pickled tries are loadable from FastAPI without __main__-shenanigans. """ from __future__ import annotations import csv import json import pickle from dataclasses import dataclass, asdict from pathlib import Path from typing import Optional UNKNOWN_TOKEN = 'Unknown' @dataclass class SourceRecord: source: str # "dickson" | "bbaw" | "tla" translit: str translation: str freq: int = 1 # The canonical Gardiner sequence as it was in the source corpus, # BEFORE any OOV->Unknown substitution. Useful when several real # sequences collapse onto the same trie path through Unknown. gardiner_orig: Optional[str] = None class TrieNode: __slots__ = ('children', 'is_end', 'records') def __init__(self): self.children: dict[str, 'TrieNode'] = {} self.is_end: bool = False self.records: list[SourceRecord] = [] class SphinxTrie: """Gardiner-sequence prefix trie.""" def __init__(self, allowed_codes: Optional[set[str]] = None): self.root = TrieNode() self.total_entries = 0 self.allowed_codes = set(allowed_codes) if allowed_codes else None # ── Insertion ────────────────────────────────────────────────────── def insert(self, codes: list[str], record: SourceRecord) -> bool: if not codes: return False if self.allowed_codes is not None: if any(c not in self.allowed_codes for c in codes): return False node = self.root for c in codes: if c not in node.children: node.children[c] = TrieNode() node = node.children[c] if not node.is_end: self.total_entries += 1 node.is_end = True node.records.append(record) return True # ── Lookup ───────────────────────────────────────────────────────── def _walk(self, codes: list[str]) -> Optional[TrieNode]: node = self.root for c in codes: child = node.children.get(c) if child is None: return None node = child return node def search(self, codes: list[str]) -> Optional[list[SourceRecord]]: node = self._walk(codes) return list(node.records) if (node and node.is_end) else None def starts_with(self, codes: list[str]) -> bool: return self._walk(codes) is not None def autocomplete(self, prefix: list[str], max_results: int = 10) -> list[dict]: node = self._walk(prefix) if node is None: return [] out: list[dict] = [] self._dfs(node, list(prefix), out, max_results) return out def _dfs(self, node, path, out, max_results): if len(out) >= max_results: return if node.is_end: out.append({ 'gardiner_seq': ' '.join(path), 'records' : [asdict(r) for r in node.records], }) for c, child in node.children.items(): self._dfs(child, path + [c], out, max_results) if len(out) >= max_results: return # ── Bounded Levenshtein over the trie ────────────────────────────── def fuzzy_search(self, codes: list[str], max_distance: int = 2, max_results: int = 5) -> list[dict]: """ Bounded edit distance over Gardiner-code sequences, traversing the trie and pruning branches whose minimum row > max_distance. Ranking: distance asc, total record freq desc. """ n = len(codes) if n == 0: return [] candidates: list[tuple[int, int, list[str], TrieNode]] = [] first_row = list(range(n + 1)) for c, child in self.root.children.items(): self._fuzzy_walk(child, c, codes, first_row, [c], max_distance, candidates) candidates.sort(key=lambda x: (x[0], -x[1])) return [ { 'gardiner_seq': ' '.join(path), 'distance' : dist, 'freq' : freq, 'records' : list(node.records), # SourceRecord instances } for dist, freq, path, node in candidates[:max_results] ] def _fuzzy_walk(self, node, edge, target, prev_row, path, max_dist, candidates): n = len(target) cur_row = [prev_row[0] + 1] for i in range(1, n + 1): ins = cur_row[i - 1] + 1 dele = prev_row[i] + 1 sub = prev_row[i - 1] + (0 if target[i - 1] == edge else 1) cur_row.append(min(ins, dele, sub)) if node.is_end and cur_row[-1] <= max_dist: freq = sum(r.freq for r in node.records) candidates.append((cur_row[-1], freq, list(path), node)) if min(cur_row) <= max_dist: for c, child in node.children.items(): self._fuzzy_walk(child, c, target, cur_row, path + [c], max_dist, candidates) # ── Build from CSV ───────────────────────────────────────────────── def _ingest_row(self, row, source, oov_strategy, gardiner_col, translit_col, translation_col, freq_col, counters): """ Process one row dict. Mutates `counters` and inserts on success. """ seq_str = (row.get(gardiner_col) or '').strip() if not seq_str: counters['empty_skipped'] += 1 return codes = seq_str.split() gardiner_orig = None if self.allowed_codes is not None: has_oov = any(c not in self.allowed_codes for c in codes) if has_oov: if oov_strategy == 'skip': counters['oov_dropped'] += 1 return gardiner_orig = ' '.join(codes) new_codes = [] for c in codes: if c in self.allowed_codes: new_codes.append(c) else: new_codes.append(UNKNOWN_TOKEN) counters['codes_subbed'] += 1 codes = new_codes counters['oov_substituted'] += 1 translit = (row.get(translit_col) or '').strip() if translit_col else '' translation = (row.get(translation_col) or '').strip() if translation_col else '' if freq_col and row.get(freq_col): try: freq = int(row[freq_col]) except (ValueError, TypeError): freq = 1 else: freq = 1 rec = SourceRecord( source=source, translit=translit, translation=translation, freq=freq, gardiner_orig=gardiner_orig, ) if self.insert(codes, rec): counters['inserted'] += 1 else: counters['empty_skipped'] += 1 def build_from_csv(self, csv_path: str | Path, source: str, gardiner_col: str = 'gardiner_seq', translit_col: str = 'translit', translation_col: str = 'translation', freq_col: Optional[str] = None, oov_strategy: str = 'skip') -> dict: """ Insert every row of csv_path into the trie. Returns build stats. oov_strategy 'skip' : drop the entire row if any code is outside allowed_codes (default). 'unknown' : replace each OOV code with the literal 'Unknown' token; row is still inserted. The original sequence is preserved on the record as record.gardiner_orig. """ if oov_strategy not in ('skip', 'unknown'): raise ValueError(f"oov_strategy must be 'skip' or 'unknown', got {oov_strategy!r}") counters = {'inserted': 0, 'oov_dropped': 0, 'oov_substituted': 0, 'codes_subbed': 0, 'empty_skipped': 0} path = Path(csv_path) with open(path, newline='', encoding='utf-8') as f: for row in csv.DictReader(f): self._ingest_row(row, source, oov_strategy, gardiner_col, translit_col, translation_col, freq_col, counters) return {'csv': str(path), 'source': source, 'oov_strategy': oov_strategy, **counters, 'total_entries': self.total_entries} def build_from_dataframe(self, df, source: str, gardiner_col: str = 'gardiner_seq', translit_col: Optional[str] = 'translit', translation_col: Optional[str] = 'translation', freq_col: Optional[str] = None, oov_strategy: str = 'skip') -> dict: """ Same as build_from_csv but takes a pandas DataFrame. translit_col / translation_col may be None when the source has no such column. """ if oov_strategy not in ('skip', 'unknown'): raise ValueError(f"oov_strategy must be 'skip' or 'unknown', got {oov_strategy!r}") counters = {'inserted': 0, 'oov_dropped': 0, 'oov_substituted': 0, 'codes_subbed': 0, 'empty_skipped': 0} wanted_cols = [c for c in (gardiner_col, translit_col, translation_col, freq_col) if c] records = df[wanted_cols].to_dict(orient='records') if wanted_cols else df.to_dict(orient='records') for row in records: self._ingest_row(row, source, oov_strategy, gardiner_col, translit_col, translation_col, freq_col, counters) return {'dataframe_rows': len(df), 'source': source, 'oov_strategy': oov_strategy, **counters, 'total_entries': self.total_entries} # ── Persistence ──────────────────────────────────────────────────── def to_pickle(self, path: str | Path) -> None: with open(path, 'wb') as f: pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL) @classmethod def from_pickle(cls, path: str | Path) -> 'SphinxTrie': with open(path, 'rb') as f: obj = pickle.load(f) if not isinstance(obj, cls): raise TypeError(f'Not a SphinxTrie: {type(obj)}') return obj def to_json(self, path: str | Path) -> None: """For human inspection only — not for runtime. Bumps recursion limit since some real corpora (BBAW passages) produce trie paths well past Python's default 1000. """ import sys as _sys def enc(node): return { 'is_end' : node.is_end, 'records': [asdict(r) for r in node.records], 'children': {k: enc(v) for k, v in node.children.items()}, } data = { 'total_entries': self.total_entries, 'allowed_codes': sorted(self.allowed_codes) if self.allowed_codes else None, 'root': enc(self.root), } prev_limit = _sys.getrecursionlimit() # need ~3x max_depth headroom for json's iterencode internals _sys.setrecursionlimit(max(prev_limit, 5000)) try: with open(path, 'w', encoding='utf-8') as f: json.dump(data, f, ensure_ascii=False, indent=2) finally: _sys.setrecursionlimit(prev_limit) def to_summary_json(self, path: str | Path, sample_per_source: int = 20) -> None: """ Compact, depth-bounded inspection artifact: stats plus a small sample of entries per source. Use this instead of to_json() when the trie is large. """ n_nodes = max_depth = 0 per_source: dict[str, int] = {} samples: dict[str, list] = {} terminal_paths_by_depth: dict[int, int] = {} def visit(node, path): nonlocal n_nodes, max_depth n_nodes += 1 d = len(path) if d > max_depth: max_depth = d if node.is_end: terminal_paths_by_depth[d] = terminal_paths_by_depth.get(d, 0) + 1 for r in node.records: per_source[r.source] = per_source.get(r.source, 0) + 1 bucket = samples.setdefault(r.source, []) if len(bucket) < sample_per_source: bucket.append({ 'gardiner_seq': ' '.join(path), **asdict(r), }) for c, child in node.children.items(): visit(child, path + [c]) visit(self.root, []) # depth histogram (binned) bins = [(1, 5), (6, 10), (11, 20), (21, 30), (31, 50), (51, 100), (101, 200), (201, 1000)] depth_hist = [] for lo, hi in bins: cnt = sum(v for d, v in terminal_paths_by_depth.items() if lo <= d <= hi) depth_hist.append({'range': f'{lo}..{hi}', 'count': cnt}) data = { 'total_entries' : self.total_entries, 'total_nodes' : n_nodes, 'max_depth' : max_depth, 'records_per_source': per_source, 'terminal_depth_histogram': depth_hist, 'allowed_codes' : sorted(self.allowed_codes) if self.allowed_codes else None, 'samples_per_source': samples, } with open(path, 'w', encoding='utf-8') as f: json.dump(data, f, ensure_ascii=False, indent=2) # ── Stats / introspection ────────────────────────────────────────── def stats(self) -> dict: n_nodes = 0 max_depth = 0 per_source: dict[str, int] = {} def visit(node, d): nonlocal n_nodes, max_depth n_nodes += 1 if d > max_depth: max_depth = d for r in node.records: per_source[r.source] = per_source.get(r.source, 0) + 1 for child in node.children.values(): visit(child, d + 1) visit(self.root, 0) return { 'total_entries' : self.total_entries, 'total_nodes' : n_nodes, 'max_depth' : max_depth, 'records_per_source': per_source, } def __repr__(self) -> str: return (f'SphinxTrie(entries={self.total_entries}, ' f'allowed={len(self.allowed_codes) if self.allowed_codes else "any"})')