code/src/enhancements/broken_predict.py

#!/usr/bin/env python3
"""
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."""
        # p(word | left) × prod p(right_i | left + [word] + right[:i])
        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:
            # Only consider words appearing at least twice (speed + quality)
            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.
    """
    # Flatten all signs with global index (respect line breaks as <EOL> marker)
    flat = []
    sign_refs = []  # parallel list of (line_idx, sign_idx) or None for EOL
    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))
        # Add end-of-line marker (None = boundary, not predicted)
        flat.append("</s>")
        sign_refs.append(None)

    # For each hole, gather context
    filled = 0
    for i, tok in enumerate(flat):
        if tok is not None and tok != "</s>":
            continue
        if sign_refs[i] is None:
            continue  # EOL, skip
        # Gather left-ctx (skip holes/EOLs)
        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'  # new category
            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)}")

    # Pre-compute candidate list once
    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()