hitit-cuneiform-ocr / code /src /enhancements /category_extrapolation.py
savastakan's picture
Initial upload: code + 5 record checkpoints + fuse
f211247 verified
Raw
History Blame Contribute Delete
4.21 kB
#!/usr/bin/env python3
"""Category Extrapolation (Zhao et al., CVPR 2025) for long-tail cuneiform.
For each target ABZ sign, ask an LLM to list visually-similar sibling signs
and provide discriminating feature phrases. Output used as:
1. Hard-negative pairs for contrastive loss
2. Sibling-aware confusion graph (cluster hints)
3. Prompt for VLM re-ranking
Does NOT train or modify models — produces JSON sibling map.
"""
import os, sys, json, argparse, time
from pathlib import Path
from collections import Counter
ROOT = Path("/arf/scratch/stakan/hitit-proje")
def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)
PROMPT = """You are an expert on Hittite cuneiform signs (Borger ABZ list).
Given the sign **{target}**, list 3–5 other ABZ signs that are most easily
confused with it due to visual similarity in wedge layout (not phonetic).
Respond in this exact JSON format:
{{
"target": "{target}",
"siblings": ["SIGN_A", "SIGN_B", "SIGN_C"],
"distinguishing_features": "<one sentence describing what visually sets {target} apart>"
}}
Only ABZ sign names. No commentary outside JSON. If {target} is ambiguous or unknown, return empty siblings.
"""
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--labels-source', required=True,
help='train ckpt with label_to_idx, OR manifest jsonl')
ap.add_argument('--abz-map', default=str(ROOT / 'hitit_ocr/data/abz_to_sign.json'))
ap.add_argument('--output', required=True)
ap.add_argument('--model', default='claude-haiku-4-5-20251001')
ap.add_argument('--max-labels', type=int, default=0, help='0=all')
ap.add_argument('--delay-ms', type=int, default=0)
args = ap.parse_args()
try: import anthropic
except ImportError: log("anthropic missing"); sys.exit(1)
if 'ANTHROPIC_API_KEY' not in os.environ:
log("ANTHROPIC_API_KEY not set"); sys.exit(1)
client = anthropic.Anthropic()
# Load labels
labels = []
p = Path(args.labels_source)
if p.suffix == '.pt':
import torch
ck = torch.load(str(p), map_location='cpu', weights_only=False)
labels = list(ck['label_to_idx'].keys())
else:
seen = set()
for line in open(p):
r = json.loads(line)
if r.get('unified_label') and r['unified_label'] not in seen:
seen.add(r['unified_label']); labels.append(r['unified_label'])
labels = sorted(set(labels))
if args.max_labels: labels = labels[:args.max_labels]
log(f"Labels: {len(labels)}")
# Resume
out = {'entries': {}, 'model': args.model}
if Path(args.output).exists():
out = json.load(open(args.output))
out.setdefault('entries', {})
log(f"Resume: {len(out['entries'])} done")
for i, lab in enumerate(labels):
if lab in out['entries']: continue
try:
msg = client.messages.create(
model=args.model, max_tokens=300,
messages=[{'role': 'user',
'content': PROMPT.format(target=lab)}])
txt = msg.content[0].text.strip()
# Extract JSON
jstart = txt.find('{'); jend = txt.rfind('}')
entry = json.loads(txt[jstart:jend+1])
out['entries'][lab] = entry
except Exception as e:
out['entries'][lab] = {'target': lab, 'siblings': [],
'error': str(e)[:200]}
if (i + 1) % 20 == 0:
log(f" {i+1}/{len(labels)}")
json.dump(out, open(args.output, 'w'), indent=2)
if args.delay_ms: time.sleep(args.delay_ms / 1000.0)
# Build a siblings-only map for easy lookup in training
sib_map = {k: v.get('siblings', []) for k, v in out['entries'].items()}
out['sibling_map'] = sib_map
# Feature matrix: label → indices of siblings within known label set
label_set = set(labels)
out['sibling_edges'] = [[a, b] for a, sibs in sib_map.items()
for b in sibs if b in label_set and b != a]
log(f"Sibling edges: {len(out['sibling_edges'])}")
json.dump(out, open(args.output, 'w'), indent=2)
if __name__ == '__main__':
main()