| |
| """ |
| Predict broken / lost signs from surrounding context using KN 5-gram LM. |
| |
| Strategy: |
| For each position marked damage='broken' (lost) or damage='partial' with very |
| low confidence, gather (left-4, right-4) context signs (non-broken). |
| Score every vocab token by: |
| p_fill(w) = p_LM(w | left context) × p_LM(right context | w, left context) |
| Pick top-k; if top1 prob > threshold, use it (with ⸢⸣ half-brackets). |
| Else keep "x" (lost). |
| |
| This uses only the LM — no image model. For stronger predictions you'd add the |
| classifier's full prob distribution over the broken crop (if any crop exists), |
| and blend with LM score. |
| """ |
| import argparse |
| import json |
| import math |
| import pickle |
| from collections import defaultdict |
| from pathlib import Path |
|
|
|
|
| def load_lm(path): |
| from collections import Counter |
| with open(path, 'rb') as f: |
| d = pickle.load(f) |
| counts = {k: Counter(v) for k, v in d['counts'].items()} |
| ctypes = {k: defaultdict(set, {kk: set(vv) for kk, vv in v.items()}) |
| for k, v in d['context_types'].items()} |
| return { |
| 'counts': counts, 'ctypes': ctypes, |
| 'vocab': d['vocab'], 'order': d['order'], |
| 'D': d['discount'], |
| } |
|
|
|
|
| class LM: |
| def __init__(self, lm): |
| self.__dict__.update(lm) |
| self.ctx_total = {n: defaultdict(int) for n in range(1, self.order + 1)} |
| for n, cnt in self.counts.items(): |
| for ng, c in cnt.items(): |
| if n > 1: |
| self.ctx_total[n][ng[:-1]] += c |
| self.cont_count = defaultdict(int) |
| for ng in self.counts[2]: |
| self.cont_count[ng[-1]] += 1 |
| self.total_cont = sum(self.cont_count.values()) |
|
|
| def cont_prob(self, w): |
| return (self.cont_count.get(w, 0) + 1) / (self.total_cont + len(self.vocab) + 1) |
|
|
| def prob(self, context, word): |
| n = min(len(context) + 1, self.order) |
| while n > 1: |
| ctx = tuple(context[-(n - 1):]) |
| if self.ctx_total[n].get(ctx, 0) > 0: |
| break |
| n -= 1 |
| if n == 1: |
| return self.cont_prob(word) |
| ctx = tuple(context[-(n - 1):]) |
| c_ng = self.counts[n].get(ctx + (word,), 0) |
| c_ctx = self.ctx_total[n][ctx] |
| types = len(self.ctypes[n - 1].get(ctx, ())) if (n - 1) in self.ctypes else 0 |
| alpha = max(c_ng - self.D, 0) / c_ctx if c_ctx else 0 |
| lam = (self.D * types) / c_ctx if c_ctx else 1.0 |
| lower = self.prob(context[-(n - 2):] if n > 2 else (), word) |
| return alpha + lam * lower |
|
|
| def score_fill(self, left_ctx, right_ctx, word): |
| """Joint prob of placing `word` between left_ctx and right_ctx.""" |
| |
| lp = math.log(max(self.prob(tuple(left_ctx), word), 1e-12)) |
| seq = list(left_ctx) + [word] |
| for rw in right_ctx: |
| lp += math.log(max(self.prob(tuple(seq), rw), 1e-12)) |
| seq.append(rw) |
| return lp |
|
|
| def topk_fill(self, left_ctx, right_ctx, k=5, candidates=None): |
| """Return top-k (word, logprob) for filling a hole between contexts.""" |
| if candidates is None: |
| |
| candidates = [w for w, c in self.vocab.items() if c >= 2] |
| scores = [] |
| for w in candidates: |
| lp = self.score_fill(left_ctx, right_ctx, w) |
| scores.append((w, lp)) |
| scores.sort(key=lambda x: -x[1]) |
| return scores[:k] |
|
|
|
|
| def fill_tablet(tablet, lm, conf_thresh=math.log(0.15), ctx=4, topk=3): |
| """Walk tablet lines; for each broken/x sign, predict fill from context. |
| |
| Modifies tablet in-place, adds 'predicted_fill' field to broken signs. |
| """ |
| |
| flat = [] |
| sign_refs = [] |
| for li, L in enumerate(tablet['lines']): |
| for si, s in enumerate(L['signs']): |
| tok = s['sign'] if s.get('damage') != 'broken' else None |
| flat.append(tok) |
| sign_refs.append((li, si)) |
| |
| flat.append("</s>") |
| sign_refs.append(None) |
|
|
| |
| filled = 0 |
| for i, tok in enumerate(flat): |
| if tok is not None and tok != "</s>": |
| continue |
| if sign_refs[i] is None: |
| continue |
| |
| left = [] |
| j = i - 1 |
| while j >= 0 and len(left) < ctx: |
| if flat[j] is not None and flat[j] != "</s>": |
| left.insert(0, flat[j]) |
| j -= 1 |
| right = [] |
| j = i + 1 |
| while j < len(flat) and len(right) < ctx: |
| if flat[j] is not None and flat[j] != "</s>": |
| right.append(flat[j]) |
| j += 1 |
| topk_preds = lm.topk_fill(left, right, k=topk) |
| li, si = sign_refs[i] |
| sign = tablet['lines'][li]['signs'][si] |
| sign['predicted_fill'] = [ |
| {'word': w, 'logprob': round(lp, 3), 'prob': round(math.exp(lp / max(1, len(left) + len(right) + 1)), 4)} |
| for w, lp in topk_preds |
| ] |
| if topk_preds and topk_preds[0][1] > conf_thresh * (len(left) + len(right) + 1): |
| sign['fill_best'] = topk_preds[0][0] |
| sign['sign_original'] = sign['sign'] |
| sign['sign'] = topk_preds[0][0] |
| sign['damage'] = 'predicted' |
| filled += 1 |
| return filled |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--lm', required=True, help='pickled KN LM') |
| ap.add_argument('--structure', required=True, help='tablet_structure.jsonl') |
| ap.add_argument('--output', required=True) |
| ap.add_argument('--conf-thresh', type=float, default=0.15) |
| ap.add_argument('--ctx', type=int, default=4) |
| ap.add_argument('--topk', type=int, default=3) |
| ap.add_argument('--limit', type=int, default=0) |
| args = ap.parse_args() |
|
|
| print(f"[load] LM {args.lm}") |
| lm_data = load_lm(args.lm) |
| lm = LM(lm_data) |
| print(f" vocab={len(lm.vocab)}") |
|
|
| |
| candidates = [w for w, c in lm.vocab.items() if c >= 2] |
| print(f" candidates={len(candidates)}") |
|
|
| n = 0; total_filled = 0 |
| with open(args.structure) as f, open(args.output, 'w') as out: |
| for line in f: |
| t = json.loads(line) |
| filled = fill_tablet(t, lm, conf_thresh=math.log(args.conf_thresh), |
| ctx=args.ctx, topk=args.topk) |
| total_filled += filled |
| out.write(json.dumps(t, ensure_ascii=False) + "\n") |
| n += 1 |
| if n % 10 == 0: |
| print(f" [{n}] total_filled={total_filled}") |
| if args.limit and n >= args.limit: |
| break |
| print(f"\nDONE: {n} tablets, {total_filled} broken signs predicted → {args.output}") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
