OCR-Viewer (#1)
Co-authored-by: nathan <nathan.kehler@gmail.com> Reviewed-on: #1
This commit is contained in:
33
scripts/check_eod.py
Normal file
33
scripts/check_eod.py
Normal file
@@ -0,0 +1,33 @@
|
||||
import json, pathlib
|
||||
from collections import Counter, defaultdict
|
||||
|
||||
data = json.loads(pathlib.Path('outputs/Calgary-Highlanders_Sep44_positions.json').read_text())
|
||||
|
||||
cats = Counter(p['category'] for p in data)
|
||||
print('=== CATEGORIES ===')
|
||||
for k, v in sorted(cats.items()):
|
||||
print(f' {k}: {v}')
|
||||
|
||||
eod = [p for p in data if p['is_end_of_day']]
|
||||
print(f'\n=== EOD entries: {len(eod)} ===')
|
||||
for p in eod:
|
||||
print(f" {str(p['date']):<16} cat={p['category']:<14} grid={str(p['grid']):<8} place={p['place_name']}")
|
||||
|
||||
eod_by_date = defaultdict(list)
|
||||
for p in eod:
|
||||
eod_by_date[p['date']].append(p)
|
||||
multi = {d: ps for d, ps in eod_by_date.items() if len(ps) > 1}
|
||||
if multi:
|
||||
print('\nWARNING - multiple EOD on same date:')
|
||||
for d, ps in multi.items():
|
||||
print(f' {d}: {len(ps)} entries')
|
||||
else:
|
||||
print('\nOK: exactly one EOD per date')
|
||||
|
||||
# List dates with no EOD
|
||||
all_dates = sorted({p['date'] for p in data if p['date']})
|
||||
eod_dates = set(eod_by_date.keys())
|
||||
missing = [d for d in all_dates if d not in eod_dates]
|
||||
if missing:
|
||||
print(f'\nDates with NO EOD: {missing}')
|
||||
|
||||
723
scripts/extract_positions_old.py
Normal file
723
scripts/extract_positions_old.py
Normal file
@@ -0,0 +1,723 @@
|
||||
"""
|
||||
Extract positional data from Calgary Highlanders War Diary Sep 44 (pages 7-57).
|
||||
Outputs a JSON array of position objects.
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
INPUT_FILE = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_olmocr.md"
|
||||
OUTPUT_FILE = _PROJECT_ROOT / "outputs" / "Calgary-Highlanders_Sep44_positions.json"
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def expand_grid(raw: str) -> tuple[str, bool]:
|
||||
"""
|
||||
Return (6-figure-string, inferred).
|
||||
4-figure AABB → centre of 1km square = AA5 BB5 (grid_inferred=True)
|
||||
6-figure AAABBB → returned as-is (grid_inferred=False)
|
||||
8-figure AAAABBBB → truncate to AAA BBB (grid_inferred=False)
|
||||
"""
|
||||
s = re.sub(r'[^0-9]', '', raw)
|
||||
if len(s) == 4:
|
||||
# e.g. "8450" → easting 84, northing 50 → centre 845, 505
|
||||
e = s[0:2] + '5'
|
||||
n = s[2:4] + '5'
|
||||
return e + n, True
|
||||
if len(s) == 6:
|
||||
return s, False
|
||||
if len(s) == 8:
|
||||
# AAAABBBB: easting = AAAA (take first 3), northing = BBBB (take first 3)
|
||||
# e.g. 15618050 → E=1561 → 156, N=8050 → 805 → "156805"
|
||||
e = s[0:3]
|
||||
n = s[4:7]
|
||||
return e + n, False
|
||||
return s, False
|
||||
|
||||
|
||||
_GRID_RE = re.compile(
|
||||
r'\b'
|
||||
r'(?:MR\s*|GR\s*)?' # optional prefix
|
||||
r'([0-9]{3,4}\s*[0-9]{3,4})' # 6 or 4+4 digit grid
|
||||
r'\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
|
||||
def find_grids(text: str):
|
||||
"""Return list of (raw_match, cleaned_digits)."""
|
||||
found = []
|
||||
for m in _GRID_RE.finditer(text):
|
||||
raw = m.group(1)
|
||||
digits = re.sub(r'\s', '', raw)
|
||||
if len(digits) in (4, 6, 8):
|
||||
found.append((m.group(0).strip(), digits))
|
||||
# also pick up standalone 6-digit runs not already caught
|
||||
for m in re.finditer(r'\b([0-9]{6})\b', text):
|
||||
digits = m.group(1)
|
||||
already = any(d == digits for _, d in found)
|
||||
if not already:
|
||||
found.append((digits, digits))
|
||||
return found
|
||||
|
||||
|
||||
# ── HTML strip ────────────────────────────────────────────────────────────────
|
||||
|
||||
class _Stripper(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.parts = []
|
||||
def handle_data(self, data):
|
||||
self.parts.append(data)
|
||||
def get_text(self):
|
||||
return ' '.join(p for p in self.parts if p.strip())
|
||||
|
||||
def strip_html(html_str: str) -> str:
|
||||
s = _Stripper()
|
||||
# replace <br> with spaces for readability
|
||||
html_str = re.sub(r'<br\s*/?>', ' ', html_str, flags=re.IGNORECASE)
|
||||
s.feed(html_str)
|
||||
return s.get_text()
|
||||
|
||||
|
||||
# ── categorise ───────────────────────────────────────────────────────────────
|
||||
|
||||
_SUBUNIT_RE = re.compile(
|
||||
r'\b(Able\s+Coy|Baker\s+Coy|Charlie\s+Coy|Dog\s+Coy|'
|
||||
r'"A"\s+Coy|"B"\s+Coy|"C"\s+Coy|"D"\s+Coy|'
|
||||
r"'A'\s+Coy|'B'\s+Coy|'C'\s+Coy|'D'\s+Coy|'Able'\s+Coy|'Baker'\s+Coy|'Charlie'\s+Coy|'Dog'\s+Coy|"
|
||||
r'Carrier\s+Platoon|carriers|Pioneer\s+Platoon|pioneers|scouts?|Scout\s+Platoon|'
|
||||
r'Support\s+Coy|Anti[-\s]?[Tt]ank\s+Pl(?:atoon)?|mortar\s+platoon|'
|
||||
r'\d+\s*Pl(?:atoon)?|'
|
||||
r'\d+\s*Sec(?:tion)?)\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
_ENEMY_RE = re.compile(
|
||||
r'\b(enemy|Jerry|Hun|Boche|MG\s*pos|MMG\s*pos|sniper|'
|
||||
r'counter.attack|strongpoint|mortar\s*pos|machine\s*gun|'
|
||||
r'block\s*house|pill\s*box|88mm|S\.S\.|SS troop)\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
_PATROL_RE = re.compile(
|
||||
r'\b(patrol|recce\s*patrol|fighting\s*patrol|'
|
||||
r'standing\s*patrol|OP|O\.P\.|observation\s*post)\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
_FRIENDLY_UNITS = re.compile(
|
||||
r'\b(R\.H\.C\.|Black\s*Watch|RHC|R\s*de\s*Mais|'
|
||||
r'Fus\s*M\.R\.|F\.M\.R\.|R\.R\.C\.|'
|
||||
r'4\s*S\.S\.|Royal\s*Regt|S\.Sask|'
|
||||
r'Tor\s*Scots|Toronto\s*Scots|'
|
||||
r'White\s*Brig(?:ade)?|F\.F\.I\.|Maquis)\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
# Tac HQ and Bn HQ in narrative
|
||||
_TAC_HQ_RE = re.compile(
|
||||
r'\bTac\s*H(?:Q)?\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
_BN_HQ_RE = re.compile(
|
||||
r'\b(?:Bn\.?\s*H\.?Q\.?|Battalion\s*H\.?Q\.?|'
|
||||
r'Battle\s*H\.?Q\.?|Command\s*Post|'
|
||||
r'moved\s+(?:his\s+)?H\.?Q\.|'
|
||||
r'set\s+up\s+H\.?Q\.|'
|
||||
r'took\s+up\s+(?:a\s+)?H\.?Q\.)\b',
|
||||
re.IGNORECASE
|
||||
)
|
||||
|
||||
def categorise(sentence: str, place_name: str | None, subunit: str | None,
|
||||
friendly: str | None, is_place_col: bool = False) -> str:
|
||||
"""
|
||||
Categories (in priority order):
|
||||
TAC_HQ – Tactical HQ position
|
||||
BN_HQ – Battalion HQ position (incl. all place-column entries)
|
||||
FRIENDLY – another friendly unit's position
|
||||
ENEMY – enemy position / feature
|
||||
PATROL – patrol route/endpoint
|
||||
SUBUNIT – company or platoon position
|
||||
UNIT_MOVEMENT – general battalion move/position
|
||||
MISC – catch-all
|
||||
"""
|
||||
if is_place_col:
|
||||
# Place column always records where Bn HQ was
|
||||
if _TAC_HQ_RE.search(sentence):
|
||||
return "TAC_HQ"
|
||||
return "BN_HQ"
|
||||
if _TAC_HQ_RE.search(sentence):
|
||||
return "TAC_HQ"
|
||||
if _BN_HQ_RE.search(sentence):
|
||||
return "BN_HQ"
|
||||
if friendly:
|
||||
return "FRIENDLY"
|
||||
if _ENEMY_RE.search(sentence):
|
||||
return "ENEMY"
|
||||
if _PATROL_RE.search(sentence):
|
||||
return "PATROL"
|
||||
if subunit:
|
||||
return "SUBUNIT"
|
||||
return "UNIT_MOVEMENT"
|
||||
|
||||
|
||||
def extract_subunit(sentence: str) -> str | None:
|
||||
m = _SUBUNIT_RE.search(sentence)
|
||||
if m:
|
||||
return m.group(0).strip()
|
||||
return None
|
||||
|
||||
def extract_friendly(sentence: str) -> str | None:
|
||||
m = _FRIENDLY_UNITS.search(sentence)
|
||||
if m:
|
||||
return m.group(0).strip()
|
||||
return None
|
||||
|
||||
|
||||
# ── split sentences ──────────────────────────────────────────────────────────
|
||||
|
||||
def split_sentences(text: str):
|
||||
"""Rough sentence splitter – split on '. ' or '<br>'."""
|
||||
# normalise
|
||||
text = re.sub(r'\s+', ' ', text).strip()
|
||||
parts = re.split(r'(?<=[.!?])\s+(?=[A-Z"\'(])', text)
|
||||
return [p.strip() for p in parts if p.strip()]
|
||||
|
||||
|
||||
# ── page extractor ────────────────────────────────────────────────────────────
|
||||
|
||||
_PAGE_RE = re.compile(r'^## Page (\d+)\s*$', re.MULTILINE)
|
||||
|
||||
def extract_pages(text: str, first: int, last: int) -> str:
|
||||
pages = list(_PAGE_RE.finditer(text))
|
||||
start_idx = None
|
||||
end_idx = len(text)
|
||||
for i, m in enumerate(pages):
|
||||
n = int(m.group(1))
|
||||
if n == first and start_idx is None:
|
||||
start_idx = m.start()
|
||||
if n == last + 1 and start_idx is not None:
|
||||
end_idx = m.start()
|
||||
break
|
||||
if start_idx is None:
|
||||
return ""
|
||||
return text[start_idx:end_idx]
|
||||
|
||||
|
||||
# ── table row parser ──────────────────────────────────────────────────────────
|
||||
|
||||
_ROW_RE = re.compile(r'<tr>(.*?)</tr>', re.DOTALL | re.IGNORECASE)
|
||||
_CELL_RE = re.compile(r'<t[dh][^>]*>(.*?)</t[dh]>', re.DOTALL | re.IGNORECASE)
|
||||
|
||||
def parse_table_rows(table_html: str) -> list[dict]:
|
||||
"""Parse a single HTML table into list of {place, date, hour, summary} dicts."""
|
||||
rows = []
|
||||
for row_m in _ROW_RE.finditer(table_html):
|
||||
cells = [strip_html(c.group(1)).strip()
|
||||
for c in _CELL_RE.finditer(row_m.group(1))]
|
||||
if len(cells) < 4:
|
||||
continue
|
||||
# skip header rows
|
||||
if re.match(r'place|date|hour|summary|no\.', cells[0], re.IGNORECASE):
|
||||
continue
|
||||
place = cells[0] if cells[0] else None
|
||||
date = cells[1] if len(cells) > 1 else None
|
||||
hour = cells[2] if len(cells) > 2 else None
|
||||
summary = cells[3] if len(cells) > 3 else ""
|
||||
rows.append(dict(place=place, date=date, hour=hour, summary=summary))
|
||||
return rows
|
||||
|
||||
|
||||
# ── sheet ref extractor ───────────────────────────────────────────────────────
|
||||
_SHEET_RE = re.compile(
|
||||
# must start with a digit to avoid matching "Sheet Ste. Foy" etc.
|
||||
r'Sheet\s+(\d[A-Z0-9]*(?:\s*[&\-]\s*\d[A-Z0-9]*)?)',
|
||||
re.IGNORECASE
|
||||
)
|
||||
|
||||
def extract_sheet(place_text: str) -> str | None:
|
||||
if not place_text:
|
||||
return None
|
||||
# find the LAST occurrence (most specific sheet reference in column)
|
||||
matches = list(_SHEET_RE.finditer(place_text))
|
||||
if matches:
|
||||
return "Sheet " + matches[-1].group(1).strip()
|
||||
return None
|
||||
|
||||
# also pull MR grid from place column
|
||||
_PLACE_MR_RE = re.compile(r'MR\s*([0-9]{4,8}|\d{2,3}\s*\d{2,3})', re.IGNORECASE)
|
||||
|
||||
def extract_place_name(place_text: str) -> str | None:
|
||||
"""
|
||||
Extract the primary place name from the Place column.
|
||||
Strip country, bare MR tokens, Sheet refs, standalone digits,
|
||||
and Tac H / Fort coord suffixes.
|
||||
"""
|
||||
if not place_text:
|
||||
return None
|
||||
s = place_text
|
||||
# remove country names
|
||||
s = re.sub(r'\b(France|Belgium|Holland|Netherlands)\b', '', s, flags=re.IGNORECASE)
|
||||
# remove Sheet + value (digit-led)
|
||||
s = re.sub(r'\bSheet\s+\d[\w\s&\-]*', ' ', s, flags=re.IGNORECASE)
|
||||
# remove bare "Sheet" not followed by a digit
|
||||
s = re.sub(r'\bSheet\b', ' ', s, flags=re.IGNORECASE)
|
||||
# remove MR + optional digits/spaces
|
||||
s = re.sub(r'\bMR\b\s*[\d\s]*', ' ', s, flags=re.IGNORECASE)
|
||||
# remove Tac H + coords
|
||||
s = re.sub(r'\bTac\s*H\s*[\d]+\b', ' ', s, flags=re.IGNORECASE)
|
||||
# remove bare 4-8 digit grid refs
|
||||
s = re.sub(r'\b\d{4,8}\b', ' ', s)
|
||||
# remove "X Pub" and similar codes
|
||||
s = re.sub(r'\bX[\s\-]?Pub\b', ' ', s, flags=re.IGNORECASE)
|
||||
# remove orphan punctuation / collapse whitespace
|
||||
s = re.sub(r'[/\\|]', ' ', s)
|
||||
s = re.sub(r'\s+', ' ', s).strip().strip('.,-()')
|
||||
# If multiple tokens, take first meaningful phrase (up to the first double space separator)
|
||||
parts = [p.strip() for p in re.split(r'\s{2,}', s) if p.strip()]
|
||||
result = parts[0] if parts else s.strip()
|
||||
# Trim trailing stray words like "(outskirts...)" parenthetical noise
|
||||
result = re.sub(r'\s*\(.*\)\s*$', '', result).strip()
|
||||
return result if result else None
|
||||
|
||||
|
||||
def clean_date(raw: str) -> str | None:
|
||||
if not raw:
|
||||
return None
|
||||
d = re.sub(r'\s+', ' ', raw).strip()
|
||||
# Remove parenthetical notes like "(Cont)"
|
||||
d = re.sub(r'\(.*?\)', '', d).strip()
|
||||
# Extract DD Mon [YY] — ignore anything else in the cell (e.g. embedded grid refs)
|
||||
m = re.search(
|
||||
r'(\d{1,2})\s+(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\s*(44|1944)?',
|
||||
d, re.IGNORECASE
|
||||
)
|
||||
if m:
|
||||
day = m.group(1)
|
||||
mon = m.group(2).capitalize()
|
||||
# always use "44" as year for this diary
|
||||
return f"{day} {mon} 44"
|
||||
return None # unrecognisable — don't inherit garbage
|
||||
|
||||
|
||||
def clean_hour(raw: str) -> str | None:
|
||||
if not raw:
|
||||
return None
|
||||
raw = raw.strip()
|
||||
# Reject empty, the bare year "44", and plain digit-strings ≥4 chars (grids)
|
||||
if not raw or raw in ('0', '44', '1944'):
|
||||
return None
|
||||
if re.match(r'^\d{4,6}$', raw):
|
||||
return None # grid ref leaked into hour column, not a time
|
||||
# Extract a valid HHMM time
|
||||
m = re.search(r'\b([012]\d[0-5]\d)\b', raw)
|
||||
if m:
|
||||
return m.group(1)
|
||||
return None
|
||||
|
||||
|
||||
def hour_col_grid(raw: str) -> str | None:
|
||||
"""
|
||||
If the Hour column contains a plain 4-digit grid ref (not a time),
|
||||
return the digit string so it can be added as a BN_HQ entry.
|
||||
"""
|
||||
if not raw:
|
||||
return None
|
||||
raw = raw.strip()
|
||||
if re.match(r'^\d{4}$', raw):
|
||||
# Confirm it can't be a time (hours > 23 or minutes > 59)
|
||||
h, m2 = int(raw[:2]), int(raw[2:])
|
||||
if h > 23 or m2 > 59:
|
||||
return raw
|
||||
return None
|
||||
|
||||
|
||||
# ── place-column MR grid ──────────────────────────────────────────────────────
|
||||
|
||||
def place_col_positions(place_text: str, date: str, sheet: str) -> list[dict]:
|
||||
"""
|
||||
Extract grids from the Place column.
|
||||
Handles:
|
||||
(a) NAME + MR + GRID (e.g. 'Ste. Foy MR 2553')
|
||||
(b) NAME + bare 6-digit grid (e.g. 'Chateau Helleputte 769969')
|
||||
Returns one entry per unique grid found.
|
||||
"""
|
||||
if not place_text:
|
||||
return []
|
||||
|
||||
results = []
|
||||
seen_digits = set()
|
||||
|
||||
def _clean_label(raw: str) -> str | None:
|
||||
s = raw
|
||||
s = re.sub(r'\b(France|Belgium|Holland|Netherlands)\b', '', s, flags=re.IGNORECASE)
|
||||
s = re.sub(r'\bSheet\b[\w\s&\-]*', '', s, flags=re.IGNORECASE)
|
||||
s = re.sub(r'\bMR\b', '', s, flags=re.IGNORECASE)
|
||||
s = re.sub(r'\b\d{4,8}\b', '', s)
|
||||
s = re.sub(r'\s+', ' ', s).strip().strip('.,-()')
|
||||
# take rightmost meaningful word group
|
||||
parts = [p.strip() for p in re.split(r'\s{2,}', s) if p.strip()]
|
||||
r = parts[-1] if parts else s.strip()
|
||||
return r if r else None
|
||||
|
||||
# ── (a) NAME MR GRID ─────────────────────────────────────────────────────
|
||||
# Require digits immediately after MR (captures 4-digit and 6-digit grids)
|
||||
mr_pattern = re.compile(
|
||||
r'([A-Z][A-Za-zÀ-ÿ\s\.\-\']*?)\s+MR\s*(\d[\d\s]+\d)',
|
||||
re.IGNORECASE
|
||||
)
|
||||
for m in mr_pattern.finditer(place_text):
|
||||
label_raw = m.group(1)
|
||||
raw_digits = re.sub(r'\s', '', m.group(2))
|
||||
if len(raw_digits) not in (4, 6, 8):
|
||||
continue
|
||||
if raw_digits in seen_digits:
|
||||
continue
|
||||
seen_digits.add(raw_digits)
|
||||
|
||||
grid, inferred = expand_grid(raw_digits)
|
||||
if len(grid) != 6:
|
||||
grid = None
|
||||
|
||||
label = _clean_label(label_raw)
|
||||
# Determine category: place column = BN_HQ unless Tac H explicit
|
||||
cat = categorise(place_text, label, None, None, is_place_col=True)
|
||||
results.append(dict(
|
||||
date=date, hour=None,
|
||||
grid=grid, grid_inferred=inferred,
|
||||
place_name=label,
|
||||
sheet_ref=sheet,
|
||||
category=cat,
|
||||
subunit=None, friendly_unit=None, is_end_of_day=False,
|
||||
context=f"Place column: {label or '?'} MR {raw_digits}"
|
||||
))
|
||||
|
||||
# ── (b) NAME followed directly by bare 6-digit grid ──────────────────────
|
||||
bare_pattern = re.compile(
|
||||
r'([A-Z][A-Za-zÀ-ÿ\s\.\-\']+?)\s+(\d{6})\b'
|
||||
)
|
||||
for m in bare_pattern.finditer(place_text):
|
||||
label_raw = m.group(1)
|
||||
raw_digits = m.group(2)
|
||||
if raw_digits in seen_digits:
|
||||
continue
|
||||
# Skip if label contains only noise words
|
||||
clean = re.sub(r'\b(France|Belgium|Holland|Netherlands|Sheet|MR)\b', '',
|
||||
label_raw, flags=re.IGNORECASE).strip()
|
||||
if not clean:
|
||||
continue
|
||||
seen_digits.add(raw_digits)
|
||||
|
||||
grid, inferred = expand_grid(raw_digits)
|
||||
label = _clean_label(label_raw)
|
||||
cat = categorise(place_text, label, None, None, is_place_col=True)
|
||||
results.append(dict(
|
||||
date=date, hour=None,
|
||||
grid=grid, grid_inferred=inferred,
|
||||
place_name=label,
|
||||
sheet_ref=sheet,
|
||||
category=cat,
|
||||
subunit=None, friendly_unit=None, is_end_of_day=False,
|
||||
context=f"Place column: {label or '?'} {raw_digits}"
|
||||
))
|
||||
|
||||
# ── (c) NAME followed by bare 4-digit grid (no MR prefix) ────────────────
|
||||
bare_4_pattern = re.compile(
|
||||
r'([A-Z][A-Za-zÀ-ÿ\s\.\-\']+?)\s+(\d{4})\b'
|
||||
)
|
||||
for m in bare_4_pattern.finditer(place_text):
|
||||
label_raw = m.group(1)
|
||||
raw_digits = m.group(2)
|
||||
if raw_digits in seen_digits:
|
||||
continue
|
||||
# Only treat as grid if NOT a valid time (i.e. HH>23 or MM>59)
|
||||
h_val, mn_val = int(raw_digits[:2]), int(raw_digits[2:])
|
||||
if h_val <= 23 and mn_val <= 59:
|
||||
continue
|
||||
clean = re.sub(r'\b(France|Belgium|Holland|Netherlands|Sheet|MR)\b', '',
|
||||
label_raw, flags=re.IGNORECASE).strip()
|
||||
if not clean:
|
||||
continue
|
||||
seen_digits.add(raw_digits)
|
||||
|
||||
grid, inferred = expand_grid(raw_digits)
|
||||
label = _clean_label(label_raw)
|
||||
cat = categorise(label_raw, label, None, None, is_place_col=True)
|
||||
results.append(dict(
|
||||
date=date, hour=None,
|
||||
grid=grid if len(grid) == 6 else None,
|
||||
grid_inferred=inferred,
|
||||
place_name=label,
|
||||
sheet_ref=sheet,
|
||||
category=cat,
|
||||
subunit=None, friendly_unit=None, is_end_of_day=False,
|
||||
context=f"Place column: {label or '?'} {raw_digits}"
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ── main extraction ───────────────────────────────────────────────────────────
|
||||
|
||||
def extract_positions(md_text: str) -> list[dict]:
|
||||
section = extract_pages(md_text, 7, 57)
|
||||
|
||||
# find all tables
|
||||
table_re = re.compile(r'<table>(.*?)</table>', re.DOTALL | re.IGNORECASE)
|
||||
|
||||
all_positions = []
|
||||
last_date = None
|
||||
last_sheet = None
|
||||
|
||||
for tbl_m in table_re.finditer(section):
|
||||
tbl_html = tbl_m.group(0)
|
||||
rows = parse_table_rows(tbl_html)
|
||||
if not rows:
|
||||
continue
|
||||
|
||||
for row in rows:
|
||||
raw_date = clean_date(row.get('date', '') or '')
|
||||
raw_hour = clean_hour(row.get('hour', '') or '')
|
||||
place_text = row.get('place', '') or ''
|
||||
summary = row.get('summary', '') or ''
|
||||
|
||||
# update tracking state
|
||||
if raw_date:
|
||||
last_date = raw_date
|
||||
cur_date = last_date
|
||||
|
||||
sheet = extract_sheet(place_text) or last_sheet
|
||||
if sheet and re.search(r'\d', sheet):
|
||||
last_sheet = sheet
|
||||
|
||||
# ── place-column positions ──
|
||||
for pos in place_col_positions(place_text, cur_date, sheet):
|
||||
pos['hour'] = raw_hour
|
||||
all_positions.append(pos)
|
||||
|
||||
# ── hour-column grid (OCR sometimes puts MR ref here) ──
|
||||
hr_raw = row.get('hour', '') or ''
|
||||
hcg = hour_col_grid(hr_raw)
|
||||
if hcg:
|
||||
grid, inferred = expand_grid(hcg)
|
||||
col_place = extract_place_name(place_text)
|
||||
all_positions.append(dict(
|
||||
date=cur_date, hour=None,
|
||||
grid=grid if len(grid) == 6 else None,
|
||||
grid_inferred=inferred,
|
||||
place_name=col_place,
|
||||
sheet_ref=last_sheet,
|
||||
category="BN_HQ",
|
||||
subunit=None, friendly_unit=None, is_end_of_day=False,
|
||||
context=f"Place column (hour field): {col_place or '?'} {hcg}"
|
||||
))
|
||||
|
||||
# ── summary / narrative positions ──
|
||||
sentences = split_sentences(summary)
|
||||
# track whether this is the last sentence in the entry
|
||||
last_sentence_idx = len(sentences) - 1
|
||||
|
||||
# collect all (sentence_idx, grid_str, raw_context) tuples for this row
|
||||
row_matches = []
|
||||
for s_idx, sent in enumerate(sentences):
|
||||
grids = find_grids(sent)
|
||||
if grids:
|
||||
for raw_match, digits in grids:
|
||||
row_matches.append((s_idx, sent, digits, raw_match))
|
||||
else:
|
||||
# named positions without grids: look for capitalised place names
|
||||
named = re.findall(
|
||||
r'\b(Fme\s+\w+|Chateau\s+\w+|Fort\s+\w+|Casino|'
|
||||
r'Distillery\s+\w+|Brickworks|blockhouse|'
|
||||
r'Mardick|Dunkerque|Dunkirk|Brecht|Loon\s*Plage|'
|
||||
r'Bourbourgville|St\.\s*Folquin|Nordamsques|Montreuil|'
|
||||
r'Wommelg[ea]h?m|Antwerp|Ypres|Pasch[aeo]nd[ae]?le|'
|
||||
r'Gravelines|Le\s*Clipon|Coppenaxfort|'
|
||||
r'Lochtenberg|Eindhoven|Sternhoven|Ryckevorsel|'
|
||||
r'St\.\s*Leonard|Bindhoven|Schilde|Schelde)\b',
|
||||
sent, re.IGNORECASE
|
||||
)
|
||||
for name in named:
|
||||
row_matches.append((s_idx, sent, None, name))
|
||||
|
||||
# determine is_end_of_day: last grid-bearing sentence in last row of date?
|
||||
for i, (s_idx, sent, digits, raw_match) in enumerate(row_matches):
|
||||
is_last = (i == len(row_matches) - 1)
|
||||
|
||||
# extract hour from sentence if not in column
|
||||
hour = raw_hour
|
||||
if not hour:
|
||||
h_m = re.search(r'\b([012]\d[0-5]\d)\s*h(?:r|our|s)?', sent, re.IGNORECASE)
|
||||
if h_m:
|
||||
hour = h_m.group(1)
|
||||
else:
|
||||
h_m2 = re.search(r'\b([012]\d[0-5]\d)[A-Z]?\b', sent)
|
||||
if h_m2:
|
||||
candidate = h_m2.group(1)
|
||||
# make sure it looks like a time not a grid
|
||||
if int(candidate[:2]) <= 23 and int(candidate[2:]) <= 59:
|
||||
hour = candidate
|
||||
|
||||
# grid
|
||||
if digits:
|
||||
grid, inferred = expand_grid(digits)
|
||||
if len(grid) != 6:
|
||||
grid = None
|
||||
inferred = False
|
||||
else:
|
||||
grid = None
|
||||
inferred = False
|
||||
|
||||
# place name from sentence context
|
||||
pn_m = re.search(
|
||||
r'\b(Fme\s+\w+[\w\s]+?(?=\s+\d|\s+MR|\.|,|$)|'
|
||||
r'Chateau\s+\w+|Fort\s+\d+|Casino|Distillery\s+\w+|'
|
||||
r'Brickworks|blockhouse|moated\s+farm|'
|
||||
r'Mardick|Dunkerque|Dunkirk|Brecht|Loon\s*Plage|'
|
||||
r'Bourbourgville|St\.\s*Folquin|Nordamsques|Montreuil|'
|
||||
r'Wommelg[ea]h?m|Antwerp|Ypres|Pasch[aeo]nd[ae]?le|'
|
||||
r'Gravelines|Le\s*Clipon|Coppenaxfort|'
|
||||
r'Lochtenberg|Eindhoven|Sternhoven|Ryckevorsel|'
|
||||
r'St\.\s*Leonard|Bindhoven|Schilde|Schelde|'
|
||||
r'cross.?roads?|road\s+junction|road\s+junc\.?|'
|
||||
r'start\s+line|bridge|lock\s+gates|railway\s+st[na]|'
|
||||
r'windmill|windpump|pier|beach|canal)\b',
|
||||
sent, re.IGNORECASE)
|
||||
place_name = pn_m.group(0).strip() if pn_m else None
|
||||
|
||||
# also check the place column label
|
||||
col_pn = extract_place_name(place_text)
|
||||
if not place_name and col_pn:
|
||||
place_name = col_pn
|
||||
|
||||
subunit = extract_subunit(sent)
|
||||
friendly = extract_friendly(sent)
|
||||
category = categorise(sent, place_name, subunit, friendly, is_place_col=False)
|
||||
|
||||
# shorten context to 2 sentences max
|
||||
context = sent.strip()
|
||||
if len(context) > 300:
|
||||
context = context[:297] + "..."
|
||||
|
||||
all_positions.append(dict(
|
||||
date=cur_date,
|
||||
hour=hour,
|
||||
grid=grid,
|
||||
grid_inferred=inferred,
|
||||
place_name=place_name,
|
||||
sheet_ref=last_sheet,
|
||||
category=category,
|
||||
subunit=subunit,
|
||||
friendly_unit=friendly,
|
||||
is_end_of_day=False, # will be set in post-processing
|
||||
context=context
|
||||
))
|
||||
|
||||
return all_positions
|
||||
|
||||
|
||||
# ── EOD post-processing ───────────────────────────────────────────────────────
|
||||
|
||||
def assign_end_of_day(positions: list[dict]) -> list[dict]:
|
||||
"""
|
||||
For each date, set is_end_of_day=True on exactly ONE entry — the last
|
||||
recorded HQ position for that date (in document order).
|
||||
|
||||
Priority (highest first):
|
||||
1. Last place-column TAC_HQ with grid (context starts "Place column:")
|
||||
2. Last place-column BN_HQ with grid
|
||||
3. Last place-column BN_HQ or TAC_HQ without grid
|
||||
4. Last narrative TAC_HQ with grid
|
||||
5. Last narrative BN_HQ with grid
|
||||
6. Last UNIT_MOVEMENT with grid
|
||||
7. Last any entry with grid
|
||||
"""
|
||||
from collections import defaultdict
|
||||
|
||||
for p in positions:
|
||||
p['is_end_of_day'] = False
|
||||
|
||||
date_indices: dict[str, list[int]] = defaultdict(list)
|
||||
for i, p in enumerate(positions):
|
||||
if p['date']:
|
||||
date_indices[p['date']].append(i)
|
||||
|
||||
def _is_place_col(p):
|
||||
return (p.get('context') or '').startswith('Place column')
|
||||
|
||||
for date, indices in date_indices.items():
|
||||
def _last(cats, require_grid=True, place_col_only=False):
|
||||
matches = [
|
||||
i for i in indices
|
||||
if positions[i]['category'] in cats
|
||||
and (not require_grid or positions[i]['grid'])
|
||||
and (not place_col_only or _is_place_col(positions[i]))
|
||||
]
|
||||
return matches[-1] if matches else None
|
||||
|
||||
winner = (
|
||||
_last({'TAC_HQ'}, require_grid=True, place_col_only=True) or
|
||||
_last({'BN_HQ'}, require_grid=True, place_col_only=True) or
|
||||
_last({'TAC_HQ', 'BN_HQ'}, require_grid=False, place_col_only=True) or
|
||||
_last({'TAC_HQ'}, require_grid=True, place_col_only=False) or
|
||||
_last({'BN_HQ'}, require_grid=True, place_col_only=False) or
|
||||
_last({'UNIT_MOVEMENT'}, require_grid=True) or
|
||||
next((i for i in reversed(indices) if positions[i]['grid']), None)
|
||||
)
|
||||
|
||||
if winner is not None:
|
||||
positions[winner]['is_end_of_day'] = True
|
||||
|
||||
return positions
|
||||
|
||||
|
||||
# ── run ───────────────────────────────────────────────────────────────────────
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Reading source file …")
|
||||
text = INPUT_FILE.read_text(encoding="utf-8")
|
||||
|
||||
print("Extracting positions …")
|
||||
positions = extract_positions(text)
|
||||
|
||||
# deduplicate identical (date + grid + context[:60]) entries
|
||||
seen = set()
|
||||
unique = []
|
||||
for p in positions:
|
||||
key = (p["date"], p["grid"], p["context"][:60])
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(p)
|
||||
|
||||
print("Assigning end-of-day flags …")
|
||||
unique = assign_end_of_day(unique)
|
||||
|
||||
OUTPUT_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
OUTPUT_FILE.write_text(
|
||||
json.dumps(unique, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8"
|
||||
)
|
||||
print(f"Done. {len(unique)} positions written to {OUTPUT_FILE}")
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
298
scripts/ocr_confidence.py
Normal file
298
scripts/ocr_confidence.py
Normal file
@@ -0,0 +1,298 @@
|
||||
"""
|
||||
ocr_confidence.py
|
||||
-----------------
|
||||
Reads an existing olmOCR .md output file and scores each page for transcription
|
||||
confidence using the DeepInfra API. Does NOT re-run OCR.
|
||||
|
||||
Requirements:
|
||||
pip install requests
|
||||
|
||||
Usage:
|
||||
python ocr_confidence.py --api_key "YOUR_KEY"
|
||||
python ocr_confidence.py --api_key "YOUR_KEY" --input "path/to/file.md" --output "path/to/out.json"
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Load .env from project root (if present) — no external dependencies needed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _load_dotenv() -> None:
|
||||
env_path = Path(__file__).resolve().parent.parent / ".env"
|
||||
if not env_path.exists():
|
||||
return
|
||||
import os
|
||||
with open(env_path, encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
key, _, value = line.partition("=")
|
||||
os.environ.setdefault(key.strip(), value.strip())
|
||||
|
||||
_load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration — mirrors ocr_wardiaries.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||||
# olmOCR is a vision model — use a chat LLM for text-based confidence review
|
||||
MODEL = "google/gemma-3-27b-it"
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 5 # seconds between retries
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_INPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_olmocr.md"
|
||||
DEFAULT_OUTPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_confidence.json"
|
||||
|
||||
CONFIDENCE_PROMPT = (
|
||||
"You are reviewing an OCR transcription of a WWII war diary. Your job is to identify words that may have been misread by the OCR scanner — not to correct the soldier's original spelling or interpret abbreviations. Rate OCR accuracy only. Do not suggest what words "should" be. Respond in JSON only:
|
||||
{"score": 7, "uncertain_words": ["Loon", "Fme"], "notes": "Possible OCR misread in line 3"}
|
||||
)
|
||||
|
||||
ERROR_RESULT = {"score": 0, "uncertain_words": [], "notes": "API error"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parse olmOCR markdown into {page_num: text} dict
|
||||
# Same logic as parseOCRByPage() in p44-ocr-viewer.html
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_ocr_by_page(raw_text: str) -> dict[int, str]:
|
||||
"""Split olmOCR markdown on '## Page N' headings into a page-keyed dict."""
|
||||
pages: dict[int, str] = {}
|
||||
lines = raw_text.split("\n")
|
||||
page_num: int | None = None
|
||||
buf: list[str] = []
|
||||
|
||||
for line in lines:
|
||||
m_head = _PAGE_HEADING.match(line)
|
||||
if m_head:
|
||||
if page_num is not None:
|
||||
pages[page_num] = "\n".join(buf)
|
||||
page_num = int(m_head.group(1))
|
||||
buf = [line] # keep the heading at the top
|
||||
elif page_num is not None:
|
||||
buf.append(line)
|
||||
|
||||
if page_num is not None:
|
||||
pages[page_num] = "\n".join(buf)
|
||||
|
||||
return pages
|
||||
|
||||
|
||||
import re as _re
|
||||
_PAGE_HEADING = _re.compile(r"^## Page (\d+)\s*$")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API call
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def score_page(page_text: str, page_num: int, api_key: str) -> dict:
|
||||
"""
|
||||
Send one page's OCR text to DeepInfra for confidence scoring.
|
||||
Returns a dict with keys: score, uncertain_words, notes.
|
||||
On failure returns ERROR_RESULT.
|
||||
"""
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"{CONFIDENCE_PROMPT}\n\n"
|
||||
f"--- BEGIN OCR TEXT (page {page_num}) ---\n"
|
||||
f"{page_text}\n"
|
||||
f"--- END OCR TEXT ---"
|
||||
),
|
||||
}
|
||||
],
|
||||
"max_tokens": 256,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
DEEPINFRA_API_URL, headers=headers, json=payload, timeout=60
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw_content = response.json()["choices"][0]["message"]["content"].strip()
|
||||
|
||||
# Strip markdown code fences if the model wrapped the JSON
|
||||
if raw_content.startswith("```"):
|
||||
raw_content = _re.sub(r"^```[a-z]*\n?", "", raw_content)
|
||||
raw_content = _re.sub(r"\n?```$", "", raw_content)
|
||||
|
||||
result = json.loads(raw_content)
|
||||
|
||||
# Validate expected keys; fill missing ones with defaults
|
||||
score = int(result.get("score", 0))
|
||||
uncertain_words = result.get("uncertain_words", [])
|
||||
notes = result.get("notes", "")
|
||||
|
||||
if not isinstance(uncertain_words, list):
|
||||
uncertain_words = []
|
||||
|
||||
return {"score": score, "uncertain_words": uncertain_words, "notes": notes}
|
||||
|
||||
except requests.exceptions.HTTPError as exc:
|
||||
warnings.warn(f"Page {page_num}: HTTP error on attempt {attempt}/{MAX_RETRIES}: {exc}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
|
||||
except json.JSONDecodeError as exc:
|
||||
warnings.warn(
|
||||
f"Page {page_num}: Malformed JSON from API on attempt {attempt}/{MAX_RETRIES}: {exc}"
|
||||
)
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
|
||||
except Exception as exc: # noqa: BLE001
|
||||
warnings.warn(f"Page {page_num}: Unexpected error on attempt {attempt}/{MAX_RETRIES}: {exc}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
|
||||
warnings.warn(f"Page {page_num}: All {MAX_RETRIES} attempts failed — storing error result.")
|
||||
return dict(ERROR_RESULT) # return a fresh copy
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Summary helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def print_summary(results: dict) -> None:
|
||||
high = sum(1 for v in results.values() if v["score"] >= 8)
|
||||
med = sum(1 for v in results.values() if 5 <= v["score"] <= 7)
|
||||
low = sum(1 for v in results.values() if 1 <= v["score"] <= 4)
|
||||
err = sum(1 for v in results.values() if v["score"] == 0)
|
||||
|
||||
total = len(results)
|
||||
print("\n" + "=" * 50)
|
||||
print(f"CONFIDENCE SCORING COMPLETE — {total} pages scored")
|
||||
print("=" * 50)
|
||||
print(f" High (8-10) : {high:>4} ({high/total*100:.1f}%)" if total else " High (8-10) : 0")
|
||||
print(f" Medium (5-7) : {med:>4} ({med/total*100:.1f}%)" if total else " Medium (5-7) : 0")
|
||||
print(f" Low (1-4) : {low:>4} ({low/total*100:.1f}%)" if total else " Low (1-4) : 0")
|
||||
if err:
|
||||
print(f" API errors : {err:>4}")
|
||||
print("=" * 50)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Score OCR confidence for each page of an olmOCR .md output file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api_key", default=None,
|
||||
help="DeepInfra API key (defaults to DEEPINFRA_API_KEY env var)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input", default=str(DEFAULT_INPUT),
|
||||
help=f"Path to olmOCR .md file. Default: {DEFAULT_INPUT}"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", default=str(DEFAULT_OUTPUT),
|
||||
help=f"Path to write confidence JSON. Default: {DEFAULT_OUTPUT}"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--delay", type=float, default=0.5,
|
||||
help="Seconds to pause between API calls (default: 0.5)."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
import os
|
||||
api_key = args.api_key or os.environ.get("DEEPINFRA_API_KEY", "")
|
||||
if not api_key:
|
||||
print("ERROR: No API key provided. Use --api_key or set DEEPINFRA_API_KEY in .env", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
input_path = Path(args.input)
|
||||
output_path = Path(args.output)
|
||||
|
||||
# ── Read source file ──────────────────────────────────────────────────────
|
||||
if not input_path.exists():
|
||||
print(f"ERROR: Input file not found: {input_path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Reading OCR source: {input_path}")
|
||||
raw_text = input_path.read_text(encoding="utf-8")
|
||||
|
||||
pages = parse_ocr_by_page(raw_text)
|
||||
if not pages:
|
||||
print("ERROR: No '## Page N' headings found in the input file.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
total_pages = len(pages)
|
||||
sorted_pages = sorted(pages.keys())
|
||||
print(f"Found {total_pages} pages (Page {sorted_pages[0]} – {sorted_pages[-1]}).")
|
||||
|
||||
# ── Load existing results (resume support) ────────────────────────────────
|
||||
results: dict[str, dict] = {}
|
||||
if output_path.exists():
|
||||
try:
|
||||
results = json.loads(output_path.read_text(encoding="utf-8"))
|
||||
already_done = len(results)
|
||||
print(f"Resuming — {already_done} page(s) already scored, skipping them.")
|
||||
except (json.JSONDecodeError, OSError) as exc:
|
||||
warnings.warn(f"Could not read existing output ({exc}); starting fresh.")
|
||||
results = {}
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ── Score each page ───────────────────────────────────────────────────────
|
||||
scored_this_run = 0
|
||||
for page_num in sorted_pages:
|
||||
page_key = str(page_num)
|
||||
|
||||
if page_key in results:
|
||||
continue # already scored — skip
|
||||
|
||||
page_text = pages[page_num]
|
||||
result = score_page(page_text, page_num, api_key)
|
||||
results[page_key] = result
|
||||
scored_this_run += 1
|
||||
|
||||
# Progress line
|
||||
score_str = str(result["score"]) if result["score"] > 0 else "ERR"
|
||||
print(f" Page {page_num}/{sorted_pages[-1]} — score: {score_str}")
|
||||
|
||||
# Write after every page so a crash loses minimal work
|
||||
output_path.write_text(
|
||||
json.dumps(results, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
if scored_this_run > 0 and page_num != sorted_pages[-1]:
|
||||
time.sleep(args.delay)
|
||||
|
||||
# ── Final save & summary ──────────────────────────────────────────────────
|
||||
output_path.write_text(
|
||||
json.dumps(results, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8"
|
||||
)
|
||||
print(f"\nResults written to: {output_path}")
|
||||
print_summary(results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
462
scripts/ocr_correction_test.py
Normal file
462
scripts/ocr_correction_test.py
Normal file
@@ -0,0 +1,462 @@
|
||||
"""
|
||||
ocr_correction_test.py
|
||||
----------------------
|
||||
Tests Gemini Pro as a text-only correction pass over olmOCR output.
|
||||
No images sent — text in, structured corrections out.
|
||||
|
||||
Reads:
|
||||
- olmOCR output .md file (already generated)
|
||||
- CP Stacey Victory Campaign PDF (extracts relevant text)
|
||||
|
||||
Sends to Gemini Pro via DeepInfra:
|
||||
- Stacey context (unit + date relevant passages)
|
||||
- Standard abbreviations/terminology
|
||||
- olmOCR page text
|
||||
- Correction prompt
|
||||
|
||||
Outputs:
|
||||
- JSON file with suggested corrections per page
|
||||
- Summary markdown showing before/after for each suggestion
|
||||
|
||||
Usage:
|
||||
python scripts/ocr_correction_test.py
|
||||
|
||||
Requirements:
|
||||
pip install requests python-dotenv pypdf pdfplumber
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pdfplumber
|
||||
import requests
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ── Config ────────────────────────────────────────────────────────────────────
|
||||
|
||||
PROJECT_ROOT = Path("G:/IdeaProjects/AI-Prototype")
|
||||
|
||||
OLMOCR_OUTPUT = PROJECT_ROOT / "outputs_step1_olmocr" / "Calgary-Highlanders_War-Diary_Sep44_control-test_comparison_olmocr-2-7b.md"
|
||||
STACEY_PDF = PROJECT_ROOT / "Inputs" / "CP-Stacey_Official-History_The-Victory-Campaign.pdf"
|
||||
OUTPUT_DIR = PROJECT_ROOT / "outputs_step1_olmocr"
|
||||
|
||||
DEEPINFRA_API_KEY = os.getenv("DEEPINFRA_API_KEY")
|
||||
API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||||
MODEL = "google/gemini-2.5-pro"
|
||||
|
||||
# Pages to test — the ones with known olmOCR errors from comparison
|
||||
TEST_PAGES = [4, 5, 6, 7, 30, 32]
|
||||
|
||||
# ── Stacey extraction ─────────────────────────────────────────────────────────
|
||||
|
||||
# Keywords that indicate relevance to Calgary Highlanders Sep 44 operations
|
||||
STACEY_KEYWORDS = [
|
||||
"Calgary Highlanders",
|
||||
"Calg High",
|
||||
"5th Canadian Infantry Brigade",
|
||||
"5 CIB",
|
||||
"2nd Canadian Infantry Division",
|
||||
"Loon Plage",
|
||||
"Bourbourgville",
|
||||
"Brecht",
|
||||
"Scheldt",
|
||||
"Crerar",
|
||||
"MacLauchlan",
|
||||
"Dieppe",
|
||||
"September 1944",
|
||||
"Sep 44",
|
||||
]
|
||||
|
||||
MAX_STACEY_CHARS = 12_000 # ~3,000 tokens — enough context without blowing cost
|
||||
|
||||
|
||||
def extract_stacey_context(pdf_path: Path, keywords: list[str], max_chars: int) -> str:
|
||||
"""
|
||||
Extract relevant passages from Stacey by keyword matching.
|
||||
Returns a condensed string of the most relevant paragraphs.
|
||||
"""
|
||||
if not pdf_path.exists():
|
||||
print(f"WARNING: Stacey PDF not found at {pdf_path}")
|
||||
return "[Stacey context not available — PDF not found]"
|
||||
|
||||
print(f"Extracting Stacey context from {pdf_path.name}...")
|
||||
relevant_chunks = []
|
||||
|
||||
try:
|
||||
with pdfplumber.open(pdf_path) as pdf:
|
||||
print(f" Stacey: {len(pdf.pages)} pages")
|
||||
for i, page in enumerate(pdf.pages):
|
||||
text = page.extract_text()
|
||||
if not text:
|
||||
continue
|
||||
# Check if any keyword appears on this page
|
||||
text_lower = text.lower()
|
||||
if any(kw.lower() in text_lower for kw in keywords):
|
||||
relevant_chunks.append(f"[Stacey p.{i+1}]\n{text.strip()}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" WARNING: Error reading Stacey PDF: {e}")
|
||||
return "[Stacey context extraction failed]"
|
||||
|
||||
if not relevant_chunks:
|
||||
print(" WARNING: No relevant Stacey passages found")
|
||||
return "[No relevant Stacey passages found for this unit/period]"
|
||||
|
||||
combined = "\n\n".join(relevant_chunks)
|
||||
|
||||
# Truncate if needed
|
||||
if len(combined) > max_chars:
|
||||
combined = combined[:max_chars] + "\n\n[...Stacey context truncated for length...]"
|
||||
|
||||
print(f" Found {len(relevant_chunks)} relevant Stacey pages, {len(combined)} chars")
|
||||
return combined
|
||||
|
||||
|
||||
# ── olmOCR page extraction ────────────────────────────────────────────────────
|
||||
|
||||
def extract_olmocr_pages(md_path: Path, page_numbers: list[int]) -> dict[int, str]:
|
||||
"""
|
||||
Parse the olmOCR markdown output and extract specific page texts.
|
||||
Returns dict of {page_number: page_text}.
|
||||
"""
|
||||
if not md_path.exists():
|
||||
print(f"ERROR: olmOCR output not found at {md_path}")
|
||||
sys.exit(1)
|
||||
|
||||
content = md_path.read_text(encoding="utf-8")
|
||||
pages = {}
|
||||
|
||||
for page_num in page_numbers:
|
||||
# Match ## Page N — olmocr-2-7b header
|
||||
pattern = rf"## Page {page_num} — olmocr-2-7b\n(.*?)(?=\n## Page |\Z)"
|
||||
match = re.search(pattern, content, re.DOTALL)
|
||||
if match:
|
||||
pages[page_num] = match.group(1).strip()
|
||||
else:
|
||||
print(f" WARNING: Page {page_num} not found in olmOCR output")
|
||||
|
||||
return pages
|
||||
|
||||
|
||||
# ── Reference context ─────────────────────────────────────────────────────────
|
||||
|
||||
ABBREVIATIONS_CONTEXT = """
|
||||
STANDARD WWII CANADIAN ARMY ABBREVIATIONS AND TERMINOLOGY:
|
||||
|
||||
Unit abbreviations:
|
||||
- Bn = Battalion
|
||||
- Coy = Company
|
||||
- Bde = Brigade
|
||||
- Div = Division
|
||||
- RHC = Royal Hamilton Light Infantry (Rileys)
|
||||
- R de Mais / R de M = Regiment de Maisonneuve
|
||||
- Calg Highrs / Calgary H = Calgary Highlanders
|
||||
- 5 CIB = 5th Canadian Infantry Brigade
|
||||
- 2 CID / 2 CDN INF DIV = 2nd Canadian Infantry Division
|
||||
- RCA = Royal Canadian Artillery
|
||||
- RCASC = Royal Canadian Army Service Corps
|
||||
- RCEME = Royal Canadian Electrical and Mechanical Engineers
|
||||
|
||||
Rank abbreviations:
|
||||
- Lt.-Col. / Lieut.-Col. = Lieutenant Colonel
|
||||
- Maj. / Major = Major
|
||||
- Capt. = Captain
|
||||
- Lieut. / Lt. = Lieutenant
|
||||
- Cpl. = Corporal
|
||||
- Pte. = Private
|
||||
- C.O. = Commanding Officer
|
||||
- I.O. = Intelligence Officer
|
||||
- Adjt. / Adj. = Adjutant
|
||||
- B.M. = Brigade Major
|
||||
- O.C. = Officer Commanding
|
||||
|
||||
Military terms:
|
||||
- Bn HQ = Battalion Headquarters
|
||||
- "O" Group = Orders Group (briefing)
|
||||
- recce = reconnaissance
|
||||
- embussed = loaded into vehicles
|
||||
- debussed = unloaded from vehicles
|
||||
- MG = Machine Gun
|
||||
- pdrs = pounders (artillery shell weight, e.g. 25 pdrs = 25-pounder guns)
|
||||
- MR = Map Reference
|
||||
- Grid refs are 6-digit numbers referencing military map sheets
|
||||
- NTR = Nothing To Report
|
||||
- DF = Defensive Fire (artillery task)
|
||||
- SOS = Special On-call (artillery task, highest priority)
|
||||
|
||||
Key personnel Sep 44:
|
||||
- Lt.-Col. D.G. MacLauchlan = CO, Calgary Highlanders
|
||||
- General H.D.G. Crerar = Commander, First Canadian Army (NOT Montgomery)
|
||||
- Brigadier W.J. Megill = Commander, 5th Canadian Infantry Brigade
|
||||
- Major R.G. Slater = Brigade Major, 5 CIB
|
||||
- R.L. Ellis = Major, Calgary Highlanders
|
||||
- Major Kearns = OC Able Company, Calgary Highlanders
|
||||
|
||||
Key locations Sep 44 (Calgary Highlanders axis):
|
||||
- Dieppe, France (2-3 Sep 44) — liberation, memorial service
|
||||
- Loon Plage, France (7-8 Sep 44) — NOT "Loon Place", NOT "Caen Place"
|
||||
- Bourbourgville, France (7 Sep 44) — canal crossing area
|
||||
- Borsbeek, Belgium (16-18 Sep 44) — 5 CIB HQ
|
||||
- Antwerp, Belgium (Sep 44)
|
||||
- Brecht, Belgium (30 Sep 44) — Able Coy street fighting
|
||||
|
||||
Note: The march past at Dieppe on 3 Sep 44 was for General Crerar
|
||||
(First Canadian Army), NOT General Montgomery.
|
||||
"""
|
||||
|
||||
|
||||
# ── Correction prompt ─────────────────────────────────────────────────────────
|
||||
|
||||
def build_correction_prompt(page_num: int, ocr_text: str, stacey_context: str) -> str:
|
||||
return f"""You are a WWII military history expert helping to correct OCR errors in Canadian Army war diary transcriptions.
|
||||
|
||||
You have been given:
|
||||
1. A reference block of standard abbreviations and known facts about the Calgary Highlanders, September 1944
|
||||
2. Relevant passages from C.P. Stacey's official history "The Victory Campaign" (Vol III of the Official History of the Canadian Army in the Second World War)
|
||||
3. Raw OCR output from page {page_num} of the Calgary Highlanders War Diary, September 1944
|
||||
|
||||
Your task: identify OCR errors and suggest corrections. Focus on:
|
||||
- Wrong names (people, places, units)
|
||||
- Garbled military terms or abbreviations
|
||||
- Repeated text loops (where the OCR got stuck repeating a phrase)
|
||||
- Missing words or broken sentences
|
||||
- Factual errors that contradict the reference context
|
||||
|
||||
Do NOT correct:
|
||||
- Original spelling errors made by the diarist (these are historically important)
|
||||
- Archaic spellings or period-appropriate language
|
||||
- Abbreviations that are standard for the period
|
||||
|
||||
Return ONLY a JSON array. Each element must have exactly these fields:
|
||||
"page": {page_num},
|
||||
"original": "the exact text as it appears in the OCR output",
|
||||
"correction": "the corrected text",
|
||||
"confidence": "high" | "medium" | "low",
|
||||
"reason": "brief explanation of why this is an error"
|
||||
|
||||
If you find no errors, return an empty array: []
|
||||
Do not return any text outside the JSON array.
|
||||
|
||||
═══════════════════════════════════════════════
|
||||
REFERENCE: ABBREVIATIONS AND KNOWN FACTS
|
||||
═══════════════════════════════════════════════
|
||||
{ABBREVIATIONS_CONTEXT}
|
||||
|
||||
═══════════════════════════════════════════════
|
||||
REFERENCE: C.P. STACEY — THE VICTORY CAMPAIGN
|
||||
═══════════════════════════════════════════════
|
||||
{stacey_context}
|
||||
|
||||
═══════════════════════════════════════════════
|
||||
OCR TEXT TO CORRECT — PAGE {page_num}
|
||||
═══════════════════════════════════════════════
|
||||
{ocr_text}
|
||||
"""
|
||||
|
||||
|
||||
# ── API call ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def call_gemini_correction(prompt: str, page_num: int) -> list[dict]:
|
||||
"""
|
||||
Send correction request to Gemini Pro via DeepInfra.
|
||||
Returns list of correction dicts.
|
||||
"""
|
||||
if not DEEPINFRA_API_KEY:
|
||||
print("ERROR: DEEPINFRA_API_KEY not set in .env")
|
||||
sys.exit(1)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {DEEPINFRA_API_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"max_tokens": 2000,
|
||||
"temperature": 0.1, # Low temp — we want deterministic corrections
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
print(f" Calling Gemini Pro for page {page_num}...")
|
||||
print(f" Prompt length: {len(prompt):,} chars (~{len(prompt)//4:,} tokens)")
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except requests.exceptions.Timeout:
|
||||
print(f" ERROR: Timeout on page {page_num}")
|
||||
return []
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f" ERROR: API call failed: {e}")
|
||||
return []
|
||||
|
||||
data = response.json()
|
||||
|
||||
# Extract content
|
||||
try:
|
||||
raw = data["choices"][0]["message"]["content"].strip()
|
||||
except (KeyError, IndexError) as e:
|
||||
print(f" ERROR: Unexpected response structure: {e}")
|
||||
print(f" Response: {data}")
|
||||
return []
|
||||
|
||||
# Strip think blocks if present
|
||||
raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()
|
||||
|
||||
# Strip markdown code fences if present
|
||||
raw = re.sub(r"^```(?:json)?\s*", "", raw)
|
||||
raw = re.sub(r"\s*```$", "", raw)
|
||||
raw = raw.strip()
|
||||
|
||||
# Parse JSON
|
||||
try:
|
||||
corrections = json.loads(raw)
|
||||
if not isinstance(corrections, list):
|
||||
print(f" WARNING: Expected JSON array, got {type(corrections)}")
|
||||
return []
|
||||
print(f" Found {len(corrections)} correction(s)")
|
||||
return corrections
|
||||
except json.JSONDecodeError as e:
|
||||
print(f" ERROR: JSON parse failed: {e}")
|
||||
print(f" Raw response (first 500 chars): {raw[:500]}")
|
||||
return []
|
||||
|
||||
|
||||
# ── Output formatting ─────────────────────────────────────────────────────────
|
||||
|
||||
def format_corrections_report(all_corrections: dict[int, list[dict]]) -> str:
|
||||
"""
|
||||
Format all corrections into a readable markdown report.
|
||||
"""
|
||||
lines = [
|
||||
"# OCR Correction Test Report",
|
||||
"## Gemini 2.5 Pro text correction with Stacey context",
|
||||
"### Calgary Highlanders War Diary Sep 44 — olmOCR output",
|
||||
"",
|
||||
f"Pages tested: {sorted(all_corrections.keys())}",
|
||||
"",
|
||||
"---",
|
||||
"",
|
||||
]
|
||||
|
||||
total_corrections = sum(len(v) for v in all_corrections.values())
|
||||
lines.append(f"**Total corrections suggested: {total_corrections}**")
|
||||
lines.append("")
|
||||
|
||||
confidence_counts = {"high": 0, "medium": 0, "low": 0}
|
||||
|
||||
for page_num in sorted(all_corrections.keys()):
|
||||
corrections = all_corrections[page_num]
|
||||
lines.append(f"## Page {page_num}")
|
||||
|
||||
if not corrections:
|
||||
lines.append("*No corrections suggested.*")
|
||||
lines.append("")
|
||||
continue
|
||||
|
||||
lines.append(f"*{len(corrections)} correction(s) suggested*")
|
||||
lines.append("")
|
||||
|
||||
for i, c in enumerate(corrections, 1):
|
||||
conf = c.get("confidence", "unknown")
|
||||
confidence_counts[conf] = confidence_counts.get(conf, 0) + 1
|
||||
|
||||
conf_emoji = {"high": "🔴", "medium": "🟡", "low": "🟢"}.get(conf, "⚪")
|
||||
|
||||
lines.append(f"### Correction {i} — {conf_emoji} {conf.upper()} confidence")
|
||||
lines.append(f"**Original:** `{c.get('original', '[missing]')}`")
|
||||
lines.append(f"**Suggested:** `{c.get('correction', '[missing]')}`")
|
||||
lines.append(f"**Reason:** {c.get('reason', '[no reason given]')}")
|
||||
lines.append("")
|
||||
|
||||
lines.append("---")
|
||||
lines.append("## Confidence Summary")
|
||||
lines.append(f"- High confidence: {confidence_counts.get('high', 0)}")
|
||||
lines.append(f"- Medium confidence: {confidence_counts.get('medium', 0)}")
|
||||
lines.append(f"- Low confidence: {confidence_counts.get('low', 0)}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("OCR Correction Test — Gemini Pro + Stacey context")
|
||||
print("=" * 60)
|
||||
|
||||
# 1. Extract Stacey context (done once, reused for all pages)
|
||||
stacey_context = extract_stacey_context(STACEY_PDF, STACEY_KEYWORDS, MAX_STACEY_CHARS)
|
||||
|
||||
# 2. Extract olmOCR pages
|
||||
print(f"\nExtracting olmOCR pages {TEST_PAGES} from {OLMOCR_OUTPUT.name}...")
|
||||
ocr_pages = extract_olmocr_pages(OLMOCR_OUTPUT, TEST_PAGES)
|
||||
print(f" Extracted {len(ocr_pages)} pages")
|
||||
|
||||
# 3. Run correction pass for each page
|
||||
all_corrections = {}
|
||||
|
||||
for page_num in TEST_PAGES:
|
||||
print(f"\n--- Page {page_num} ---")
|
||||
|
||||
if page_num not in ocr_pages:
|
||||
print(f" Skipping — page not found in olmOCR output")
|
||||
all_corrections[page_num] = []
|
||||
continue
|
||||
|
||||
ocr_text = ocr_pages[page_num]
|
||||
|
||||
if "[OCR FAILED" in ocr_text:
|
||||
print(f" Skipping — olmOCR failed on this page")
|
||||
all_corrections[page_num] = []
|
||||
continue
|
||||
|
||||
prompt = build_correction_prompt(page_num, ocr_text, stacey_context)
|
||||
corrections = call_gemini_correction(prompt, page_num)
|
||||
all_corrections[page_num] = corrections
|
||||
|
||||
# 4. Save JSON output
|
||||
json_path = OUTPUT_DIR / "ocr_correction_test_results.json"
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(all_corrections, f, indent=2, ensure_ascii=False)
|
||||
print(f"\nJSON results saved to {json_path}")
|
||||
|
||||
# 5. Save markdown report
|
||||
report = format_corrections_report(all_corrections)
|
||||
md_path = OUTPUT_DIR / "ocr_correction_test_report.md"
|
||||
with open(md_path, "w", encoding="utf-8") as f:
|
||||
f.write(report)
|
||||
print(f"Markdown report saved to {md_path}")
|
||||
|
||||
# 6. Print summary to console
|
||||
print("\n" + "=" * 60)
|
||||
print("RESULTS SUMMARY")
|
||||
print("=" * 60)
|
||||
for page_num, corrections in sorted(all_corrections.items()):
|
||||
if corrections:
|
||||
print(f"\nPage {page_num}: {len(corrections)} correction(s)")
|
||||
for c in corrections:
|
||||
conf = c.get("confidence", "?").upper()
|
||||
print(f" [{conf}] '{c.get('original', '')}' → '{c.get('correction', '')}'")
|
||||
print(f" {c.get('reason', '')}")
|
||||
else:
|
||||
print(f"\nPage {page_num}: No corrections suggested")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
321
scripts/ocr_model_comparison.py
Normal file
321
scripts/ocr_model_comparison.py
Normal file
@@ -0,0 +1,321 @@
|
||||
"""
|
||||
ocr_model_comparison.py
|
||||
-----------------------
|
||||
Runs multiple vision LLMs on the same war diary pages and outputs
|
||||
labelled plain text for side-by-side accuracy comparison.
|
||||
|
||||
One output file per model, named:
|
||||
{diary_stem}_comparison_{model_label}.md
|
||||
|
||||
Each file uses the same ## Page N format as the main OCR pipeline
|
||||
so you can compare directly against the olmOCR output.
|
||||
|
||||
Requirements:
|
||||
pip install requests pillow python-dotenv
|
||||
|
||||
Poppler must be on PATH.
|
||||
API key read from DEEPINFRA_API_KEY in .env file.
|
||||
|
||||
Usage:
|
||||
python scripts/ocr_model_comparison.py ^
|
||||
--single "G:/path/to/diary.pdf" ^
|
||||
--output_dir "G:/path/to/outputs_step1_olmocr" ^
|
||||
--first_page 7 ^
|
||||
--last_page 16
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
import requests
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Models to test
|
||||
# ---------------------------------------------------------------------------
|
||||
# Add or remove models here. label is used in the output filename.
|
||||
# All use the same DeepInfra OpenAI-compatible endpoint.
|
||||
|
||||
MODELS = [
|
||||
{
|
||||
"label": "gemini-2.5-pro",
|
||||
"model": "google/gemini-2.5-pro",
|
||||
},
|
||||
{
|
||||
"label": "qwen3-vl-30b",
|
||||
"model": "Qwen/Qwen3-VL-30B-A3B-Instruct",
|
||||
},
|
||||
{
|
||||
"label": "olmocr-2-7b",
|
||||
"model": "allenai/olmOCR-2-7B-1025",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||||
DPI = 150
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 15
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt — same for all models so comparison is fair
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are transcribing a scanned WWII Canadian Army war diary page.
|
||||
Transcribe every word exactly as written. Do not correct spelling,
|
||||
grammar, or abbreviations. Do not add commentary or explanations.
|
||||
Keep grid references exactly as written (e.g. MR 2553, GR 442891).
|
||||
Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O.).
|
||||
Transcribe all text visible on the page — headers, titles, form numbers,
|
||||
stamps, handwritten notes, signatures, and margin notations. Nothing omitted.
|
||||
If the page contains a table with Place / Date / Hour / Summary columns,
|
||||
output it as an HTML table using <table>, <tr>, <th>, <td>, <br> tags,
|
||||
preserving column structure exactly.
|
||||
If the page is genuinely blank, output only: [BLANK]\
|
||||
"""
|
||||
|
||||
USER_PROMPT = (
|
||||
"Transcribe every word on this page exactly as written. "
|
||||
"Use an HTML table for any Place/Date/Hour/Summary diary table. "
|
||||
"Nothing omitted, nothing added."
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI,
|
||||
first_page: int = None, last_page: int = None) -> list[tuple[int, Path]]:
|
||||
prefix = output_dir / pdf_path.stem
|
||||
cmd = ["pdftoppm", "-r", str(dpi), "-png"]
|
||||
if first_page is not None:
|
||||
cmd += ["-f", str(first_page)]
|
||||
if last_page is not None:
|
||||
cmd += ["-l", str(last_page)]
|
||||
cmd += [str(pdf_path), str(prefix)]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"pdftoppm failed:\n{result.stderr}")
|
||||
|
||||
images = sorted(output_dir.glob(f"{pdf_path.stem}-*.png"))
|
||||
result_pairs = []
|
||||
for img_path in images:
|
||||
m = re.search(r'-(\d+)\.png$', img_path.name)
|
||||
page_num = int(m.group(1)) if m else len(result_pairs) + (first_page or 1)
|
||||
result_pairs.append((page_num, img_path))
|
||||
return result_pairs
|
||||
|
||||
|
||||
def image_to_base64(image_path: Path) -> str:
|
||||
with open(image_path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def ocr_page(image_path: Path, api_key: str, model: str, page_num: int) -> str:
|
||||
b64 = image_to_base64(image_path)
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": SYSTEM_PROMPT,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{b64}"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": USER_PROMPT,
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
DEEPINFRA_API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=180,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data["choices"][0]["message"]["content"].strip()
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — {model} — HTTP error: {e}]"
|
||||
except Exception as e:
|
||||
print(f" Error attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — {model} — Error: {e}]"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_comparison(pdf_path: Path, output_dir: Path, api_key: str,
|
||||
first_page: int = None, last_page: int = None):
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"OCR Model Comparison: {pdf_path.name}")
|
||||
if first_page or last_page:
|
||||
print(f"Pages: {first_page or 1} → {last_page or 'end'}")
|
||||
print(f"Models: {', '.join(m['label'] for m in MODELS)}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
print(f"\nRendering PDF pages at {DPI} DPI...")
|
||||
page_images = pdf_to_images(pdf_path, tmp_path,
|
||||
first_page=first_page, last_page=last_page)
|
||||
total = len(page_images)
|
||||
print(f"Found {total} pages.")
|
||||
|
||||
# Store results per model: {label: [(page_num, text), ...]}
|
||||
results = {m["label"]: [] for m in MODELS}
|
||||
|
||||
for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
|
||||
print(f"\nPage {pdf_page_num} ({i}/{total})")
|
||||
|
||||
for model_cfg in MODELS:
|
||||
label = model_cfg["label"]
|
||||
model = model_cfg["model"]
|
||||
print(f" [{label}]...", end=" ", flush=True)
|
||||
|
||||
page_text = ocr_page(image_path, api_key, model, pdf_page_num)
|
||||
results[label].append((pdf_page_num, page_text))
|
||||
print("done")
|
||||
|
||||
# Delay between models to avoid rate limits
|
||||
time.sleep(1)
|
||||
|
||||
# Slightly longer delay between pages
|
||||
if i < total:
|
||||
time.sleep(1)
|
||||
|
||||
# Write one output file per model
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stem = pdf_path.stem
|
||||
|
||||
for model_cfg in MODELS:
|
||||
label = model_cfg["label"]
|
||||
model_id = model_cfg["model"]
|
||||
output_path = output_dir / f"{stem}_comparison_{label}.md"
|
||||
|
||||
lines = [
|
||||
f"# {stem}",
|
||||
f"Model: {model_id}",
|
||||
f"Label: {label}",
|
||||
f"Pages: {first_page or 1}–{last_page or total}",
|
||||
f"DPI: {DPI}",
|
||||
"",
|
||||
]
|
||||
|
||||
for page_num, page_text in results[label]:
|
||||
lines.append(f"## Page {page_num} — {label}")
|
||||
lines.append("")
|
||||
lines.append(page_text)
|
||||
lines.append("")
|
||||
|
||||
output_path.write_text("\n".join(lines), encoding="utf-8")
|
||||
print(f"\nSaved: {output_path.name}")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Complete. {len(MODELS)} comparison files written to:")
|
||||
print(f" {output_dir}")
|
||||
print("\nFiles:")
|
||||
for model_cfg in MODELS:
|
||||
print(f" {stem}_comparison_{model_cfg['label']}.md")
|
||||
print(f"\nCompare these against:")
|
||||
print(f" {stem}_olmocr.md (existing olmOCR output)")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Compare multiple vision LLMs on the same war diary pages."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--single",
|
||||
required=True,
|
||||
help="Path to a single PDF file to process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
required=True,
|
||||
help="Folder to write comparison output files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--first_page",
|
||||
type=int,
|
||||
default=None,
|
||||
help="First page to process (1-based).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--last_page",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Last page to process (1-based).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = os.getenv("DEEPINFRA_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: DEEPINFRA_API_KEY not set in .env file")
|
||||
sys.exit(1)
|
||||
|
||||
pdf_path = Path(args.single)
|
||||
if not pdf_path.exists():
|
||||
print(f"Error: PDF not found: {pdf_path}")
|
||||
sys.exit(1)
|
||||
|
||||
run_comparison(
|
||||
pdf_path,
|
||||
Path(args.output_dir),
|
||||
api_key,
|
||||
first_page=args.first_page,
|
||||
last_page=args.last_page,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
271
scripts/ocr_model_comparison_anthropic.py
Normal file
271
scripts/ocr_model_comparison_anthropic.py
Normal file
@@ -0,0 +1,271 @@
|
||||
"""
|
||||
ocr_model_comparison_anthropic.py
|
||||
----------------------------------
|
||||
Runs Claude Sonnet 4.6 via the Anthropic API on war diary pages.
|
||||
Produces the same ## Page N — label output format as ocr_model_comparison.py
|
||||
so results can be directly compared against Gemini Pro, olmOCR, etc.
|
||||
|
||||
Output file:
|
||||
{diary_stem}_comparison_claude-sonnet-4-6.md
|
||||
|
||||
Requirements:
|
||||
pip install requests pillow python-dotenv
|
||||
|
||||
Poppler must be on PATH.
|
||||
ANTHROPIC_API_KEY read from .env file.
|
||||
|
||||
Usage:
|
||||
python scripts/ocr_model_comparison_anthropic.py ^
|
||||
--single "G:/IdeaProjects/AI-Prototype/Inputs/Calgary-Highlanders_War-Diary_Sep44.pdf" ^
|
||||
--output_dir "G:/IdeaProjects/AI-Prototype/outputs_step1_olmocr" ^
|
||||
--first_page 7 ^
|
||||
--last_page 32
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
import requests
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ── Config ────────────────────────────────────────────────────────────────────
|
||||
|
||||
ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
|
||||
MODEL_ID = "claude-sonnet-4-6"
|
||||
MODEL_LABEL = "claude-sonnet-4-6"
|
||||
DPI = 150
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 15
|
||||
|
||||
# ── Prompt — identical to ocr_model_comparison.py for fair comparison ─────────
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are transcribing a scanned WWII Canadian Army war diary page.
|
||||
Transcribe every word exactly as written. Do not correct spelling,
|
||||
grammar, or abbreviations. Do not add commentary or explanations.
|
||||
Keep grid references exactly as written (e.g. MR 2553, GR 442891).
|
||||
Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O.).
|
||||
Transcribe all text visible on the page — headers, titles, form numbers,
|
||||
stamps, handwritten notes, signatures, and margin notations. Nothing omitted.
|
||||
If the page contains a table with Place / Date / Hour / Summary columns,
|
||||
output it as an HTML table using <table>, <tr>, <th>, <td>, <br> tags,
|
||||
preserving column structure exactly.
|
||||
If the page is genuinely blank, output only: [BLANK]\
|
||||
"""
|
||||
|
||||
USER_PROMPT = (
|
||||
"Transcribe every word on this page exactly as written. "
|
||||
"Use an HTML table for any Place/Date/Hour/Summary diary table. "
|
||||
"Nothing omitted, nothing added."
|
||||
)
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI,
|
||||
first_page: int = None, last_page: int = None) -> list:
|
||||
prefix = output_dir / pdf_path.stem
|
||||
cmd = ["pdftoppm", "-r", str(dpi), "-png"]
|
||||
if first_page is not None:
|
||||
cmd += ["-f", str(first_page)]
|
||||
if last_page is not None:
|
||||
cmd += ["-l", str(last_page)]
|
||||
cmd += [str(pdf_path), str(prefix)]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"pdftoppm failed:\n{result.stderr}")
|
||||
|
||||
images = sorted(output_dir.glob(f"{pdf_path.stem}-*.png"))
|
||||
result_pairs = []
|
||||
for img_path in images:
|
||||
m = re.search(r'-(\d+)\.png$', img_path.name)
|
||||
page_num = int(m.group(1)) if m else len(result_pairs) + (first_page or 1)
|
||||
result_pairs.append((page_num, img_path))
|
||||
return result_pairs
|
||||
|
||||
|
||||
def image_to_base64(image_path: Path) -> str:
|
||||
"""
|
||||
Load PNG, compress to JPEG at 85% quality, return base64.
|
||||
Anthropic has a 5MB image limit — war diary PNGs at 150 DPI
|
||||
are ~6-7MB. JPEG at 85% brings them to ~0.8-1.2MB with no
|
||||
meaningful loss of OCR-relevant detail.
|
||||
"""
|
||||
from PIL import Image
|
||||
import io
|
||||
img = Image.open(image_path).convert("RGB")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=85)
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
|
||||
b64 = image_to_base64(image_path)
|
||||
|
||||
headers = {
|
||||
"x-api-key": api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL_ID,
|
||||
"max_tokens": 8192,
|
||||
"system": SYSTEM_PROMPT,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/jpeg",
|
||||
"data": b64,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": USER_PROMPT,
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
ANTHROPIC_API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=180,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# Print token usage
|
||||
usage = data.get("usage", {})
|
||||
input_t = usage.get("input_tokens", 0)
|
||||
output_t = usage.get("output_tokens", 0)
|
||||
cost = (input_t / 1_000_000 * 3.00) + (output_t / 1_000_000 * 15.00)
|
||||
print(f" tokens: {input_t} in / {output_t} out — ${cost:.4f}")
|
||||
|
||||
return data["content"][0]["text"].strip()
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if hasattr(e, 'response') and e.response is not None:
|
||||
print(f" Response: {e.response.text[:300]}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — {MODEL_ID} — HTTP error: {e}]"
|
||||
|
||||
except Exception as e:
|
||||
print(f" Error attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — {MODEL_ID} — Error: {e}]"
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def run_comparison(pdf_path: Path, output_dir: Path, api_key: str,
|
||||
first_page: int = None, last_page: int = None):
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"OCR — {MODEL_LABEL}")
|
||||
print(f"File: {pdf_path.name}")
|
||||
print(f"Pages: {first_page or 1} → {last_page or 'end'}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
print(f"\nRendering PDF pages at {DPI} DPI...")
|
||||
page_images = pdf_to_images(pdf_path, tmp_path,
|
||||
first_page=first_page, last_page=last_page)
|
||||
total = len(page_images)
|
||||
print(f"Found {total} pages.")
|
||||
|
||||
results = []
|
||||
|
||||
for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
|
||||
print(f"\nPage {pdf_page_num} ({i}/{total})...", end=" ", flush=True)
|
||||
page_text = ocr_page(image_path, api_key, pdf_page_num)
|
||||
results.append((pdf_page_num, page_text))
|
||||
print("done")
|
||||
|
||||
if i < total:
|
||||
time.sleep(1)
|
||||
|
||||
# Write output — same format as ocr_model_comparison.py
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stem = pdf_path.stem
|
||||
output_path = output_dir / f"{stem}_comparison_{MODEL_LABEL}.md"
|
||||
|
||||
lines = [
|
||||
f"# {stem}",
|
||||
f"Model: {MODEL_ID}",
|
||||
f"Label: {MODEL_LABEL}",
|
||||
f"Pages: {first_page or 1}–{last_page or total}",
|
||||
f"DPI: {DPI}",
|
||||
"",
|
||||
]
|
||||
|
||||
for page_num, page_text in results:
|
||||
lines.append(f"## Page {page_num} — {MODEL_LABEL}")
|
||||
lines.append("")
|
||||
lines.append(page_text)
|
||||
lines.append("")
|
||||
|
||||
output_path.write_text("\n".join(lines), encoding="utf-8")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Done. Output saved to:")
|
||||
print(f" {output_path.name}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="OCR war diary pages using Claude Sonnet 4.6 via Anthropic API."
|
||||
)
|
||||
parser.add_argument("--single", required=True, help="Path to PDF.")
|
||||
parser.add_argument("--output_dir", required=True, help="Output folder.")
|
||||
parser.add_argument("--first_page", type=int, default=None)
|
||||
parser.add_argument("--last_page", type=int, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: ANTHROPIC_API_KEY not set in .env")
|
||||
sys.exit(1)
|
||||
|
||||
pdf_path = Path(args.single)
|
||||
if not pdf_path.exists():
|
||||
print(f"Error: PDF not found: {pdf_path}")
|
||||
sys.exit(1)
|
||||
|
||||
run_comparison(
|
||||
pdf_path,
|
||||
Path(args.output_dir),
|
||||
api_key,
|
||||
first_page=args.first_page,
|
||||
last_page=args.last_page,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
370
scripts/step1_ocr_wardiaries_claude-4-6.py
Normal file
370
scripts/step1_ocr_wardiaries_claude-4-6.py
Normal file
@@ -0,0 +1,370 @@
|
||||
"""
|
||||
step1_ocr_wardiaries_claude-4-6.py
|
||||
------------------------------------
|
||||
Converts war diary PDFs to structured text using Claude Sonnet 4.6 via the Anthropic API.
|
||||
Drop-in replacement for step1_ocr_wardiaries_olmocr.py — same output format, same CLI args.
|
||||
|
||||
Output format: one .md file per PDF with ## Page N markers.
|
||||
- Diary pages with Place/Date/Hour/Summary tables are output as HTML tables,
|
||||
preserving column structure for downstream extraction.
|
||||
- All other pages (admin, appendices, forms, etc.) are output as plain text.
|
||||
|
||||
This structured output is the source of truth for the pipeline:
|
||||
- llm_extract.py reads the HTML tables for position and narrative extraction
|
||||
- step2_json_to_viewer_md.py flattens the HTML tables to readable prose for the OCR viewer
|
||||
|
||||
Requirements:
|
||||
pip install requests pillow python-dotenv
|
||||
|
||||
Poppler must be on PATH (pdftoppm command must work).
|
||||
|
||||
API key is read from ANTHROPIC_API_KEY in your .env file.
|
||||
|
||||
Usage:
|
||||
# Test run — pages 7 to 27 only:
|
||||
python scripts/step1_ocr_wardiaries_claude-4-6.py --single "G:/path/to/diary.pdf" --output_dir "G:/path/to/outputs" --first_page 7 --last_page 27
|
||||
|
||||
# Full diary:
|
||||
python scripts/step1_ocr_wardiaries_claude-4-6.py --single "G:/path/to/diary.pdf" --output_dir "G:/path/to/outputs"
|
||||
|
||||
# All PDFs in a folder:
|
||||
python scripts/step1_ocr_wardiaries_claude-4-6.py --input_dir "G:/path/to/pdfs" --output_dir "G:/path/to/outputs"
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from PIL import Image
|
||||
import requests
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
|
||||
MODEL = "claude-sonnet-4-6"
|
||||
DPI = 150
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 15
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OCR prompt — identical to step1_ocr_wardiaries_olmocr.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are transcribing a scanned WWII Canadian Army war diary page.
|
||||
|
||||
Transcribe every word exactly as written. Do not correct spelling, grammar,
|
||||
or abbreviations. Do not add commentary or explanations.
|
||||
Keep grid references exactly as written (e.g. MR 2553, GR 442891, MR 0675 Sheet 861).
|
||||
Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O., RHC, R de Mais).
|
||||
Transcribe all text visible on the page — headers, titles, form numbers,
|
||||
stamps, handwritten notes, signatures, and margin notations. Nothing omitted.
|
||||
|
||||
If the page contains a table with columns Place / Date / Hour / Summary / Remarks:
|
||||
Output the table as HTML using <table>, <tr>, <th>, <td> tags.
|
||||
- Place column: all lines exactly as written, separated by <br> tags.
|
||||
Include country names, place names, grid references, and sheet references.
|
||||
- Date column: exact date as written, <br> if split across lines.
|
||||
- Hour column: exact time as written, or empty if absent.
|
||||
- Summary column: full verbatim text, paragraphs separated by <br>. Every word. Nothing omitted.
|
||||
- Remarks column: exact text, or empty if absent.
|
||||
If the page has multiple rows (different dates or continuation entries), each row is a separate <tr>.
|
||||
Any text outside the table (headers, weather notes, stamps) is output as plain text
|
||||
before or after the <table> tag in the order it appears on the page.
|
||||
|
||||
If the page does NOT contain a Place/Date/Hour/Summary table:
|
||||
Output all text as plain text in the order it appears on the page.
|
||||
|
||||
If the page is genuinely blank, output only: [BLANK]\
|
||||
"""
|
||||
|
||||
USER_PROMPT = (
|
||||
"Transcribe every word on this page exactly as written. "
|
||||
"Use an HTML table for any Place/Date/Hour/Summary diary table. "
|
||||
"Nothing omitted, nothing added."
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI,
|
||||
first_page: int = None, last_page: int = None) -> list[tuple[int, Path]]:
|
||||
"""
|
||||
Render pages of a PDF to PNG images using pdftoppm (poppler).
|
||||
Returns a list of (pdf_page_number, image_path) tuples.
|
||||
"""
|
||||
prefix = output_dir / pdf_path.stem
|
||||
cmd = ["pdftoppm", "-r", str(dpi), "-png"]
|
||||
if first_page is not None:
|
||||
cmd += ["-f", str(first_page)]
|
||||
if last_page is not None:
|
||||
cmd += ["-l", str(last_page)]
|
||||
cmd += [str(pdf_path), str(prefix)]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"pdftoppm failed for {pdf_path.name}:\n{result.stderr}")
|
||||
|
||||
images = sorted(output_dir.glob(f"{pdf_path.stem}-*.png"))
|
||||
result_pairs = []
|
||||
for img_path in images:
|
||||
m = re.search(r'-(\d+)\.png$', img_path.name)
|
||||
page_num = int(m.group(1)) if m else len(result_pairs) + (first_page or 1)
|
||||
result_pairs.append((page_num, img_path))
|
||||
return result_pairs
|
||||
|
||||
|
||||
def image_to_base64_jpeg(image_path: Path, quality: int = 85) -> str:
|
||||
"""
|
||||
Load PNG, compress to JPEG at given quality, return base64.
|
||||
Anthropic has a 5 MB image limit — war diary PNGs at 150 DPI are ~6-7 MB.
|
||||
JPEG at 85% brings them to ~0.8-1.2 MB with no meaningful loss of OCR detail.
|
||||
"""
|
||||
img = Image.open(image_path).convert("RGB")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=quality)
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
|
||||
"""Send one page image to Claude via Anthropic API and return the text."""
|
||||
b64 = image_to_base64_jpeg(image_path)
|
||||
|
||||
headers = {
|
||||
"x-api-key": api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"max_tokens": 8192,
|
||||
"system": SYSTEM_PROMPT,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/jpeg",
|
||||
"data": b64,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": USER_PROMPT,
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
ANTHROPIC_API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=180,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# Print token usage and estimated cost
|
||||
usage = data.get("usage", {})
|
||||
input_t = usage.get("input_tokens", 0)
|
||||
output_t = usage.get("output_tokens", 0)
|
||||
cost = (input_t / 1_000_000 * 3.00) + (output_t / 1_000_000 * 15.00)
|
||||
print(f" tokens: {input_t} in / {output_t} out — ${cost:.4f}", end=" ")
|
||||
|
||||
return data["content"][0]["text"].strip()
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if hasattr(e, 'response') and e.response is not None:
|
||||
print(f" Response: {e.response.text[:300]}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — HTTP error: {e}]"
|
||||
except Exception as e:
|
||||
print(f" Error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — Error: {e}]"
|
||||
|
||||
|
||||
def has_table(page_text: str) -> bool:
|
||||
"""Return True if the page output contains an HTML table."""
|
||||
return bool(re.search(r'<table', page_text, re.IGNORECASE))
|
||||
|
||||
|
||||
def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str,
|
||||
first_page: int = None, last_page: int = None) -> Path:
|
||||
"""
|
||||
OCR a single PDF and write structured output to a .md file.
|
||||
Diary pages are output as HTML tables; all other pages as plain text.
|
||||
"""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Processing: {pdf_path.name}")
|
||||
if first_page or last_page:
|
||||
print(f"Pages: {first_page or 1} → {last_page or 'end'}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
output_md = output_dir / f"step1_{pdf_path.stem}_claude-4-6.md"
|
||||
|
||||
if output_md.exists():
|
||||
print(f" Already processed — skipping.")
|
||||
print(f" Delete {output_md.name} to reprocess.")
|
||||
return output_md
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
print(f" Rendering PDF pages to images at {DPI} DPI...")
|
||||
try:
|
||||
page_images = pdf_to_images(pdf_path, tmp_path,
|
||||
first_page=first_page, last_page=last_page)
|
||||
except RuntimeError as e:
|
||||
print(f" ERROR: {e}")
|
||||
return output_md
|
||||
|
||||
total = len(page_images)
|
||||
print(f" Found {total} pages to process.")
|
||||
|
||||
lines = [
|
||||
f"# {pdf_path.stem}",
|
||||
f"OCR by Claude Sonnet 4.6 ({MODEL}) via Anthropic",
|
||||
f"Source: {pdf_path.name} — pages {first_page or 1}–{last_page or total}",
|
||||
"",
|
||||
]
|
||||
|
||||
diary_count = 0
|
||||
other_count = 0
|
||||
|
||||
for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
|
||||
print(f" Page {pdf_page_num} ({i}/{total})...", end=" ", flush=True)
|
||||
|
||||
page_text = ocr_page(image_path, api_key, pdf_page_num)
|
||||
is_diary = has_table(page_text)
|
||||
|
||||
lines.append(f"## Page {pdf_page_num}")
|
||||
lines.append("")
|
||||
lines.append(page_text)
|
||||
lines.append("")
|
||||
|
||||
if is_diary:
|
||||
diary_count += 1
|
||||
print("done — diary (table)")
|
||||
else:
|
||||
other_count += 1
|
||||
print("done — plain text")
|
||||
|
||||
if i < total:
|
||||
time.sleep(0.5)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_md.write_text("\n".join(lines), encoding="utf-8")
|
||||
|
||||
print(f"\n Saved: {output_md}")
|
||||
print(f" Summary: {diary_count} diary pages (HTML tables), {other_count} other pages (plain text)")
|
||||
|
||||
return output_md
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="OCR war diary PDFs to structured text using Claude Sonnet 4.6 via Anthropic API."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_dir",
|
||||
default=None,
|
||||
help="Folder containing PDF files to process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
required=True,
|
||||
help="Folder to write output .md files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--single",
|
||||
default=None,
|
||||
help="Process a single PDF file instead of a whole folder.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--first_page",
|
||||
type=int,
|
||||
default=None,
|
||||
help="First PDF page to process (1-based).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--last_page",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Last PDF page to process (1-based).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: ANTHROPIC_API_KEY not set in environment or .env file")
|
||||
sys.exit(1)
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if args.single:
|
||||
pdfs = [Path(args.single)]
|
||||
elif args.input_dir:
|
||||
pdfs = sorted(Path(args.input_dir).glob("*.pdf"))
|
||||
else:
|
||||
print("Error: provide --input_dir or --single")
|
||||
sys.exit(1)
|
||||
|
||||
if not pdfs:
|
||||
print("No PDF files found.")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Found {len(pdfs)} PDF(s) to process.")
|
||||
print(f"Output directory: {output_dir}")
|
||||
print(f"Model: {MODEL}")
|
||||
print(f"DPI: {DPI}")
|
||||
|
||||
results = []
|
||||
for pdf_path in pdfs:
|
||||
output_md = ocr_pdf(
|
||||
pdf_path, output_dir, api_key,
|
||||
first_page=args.first_page,
|
||||
last_page=args.last_page,
|
||||
)
|
||||
results.append(output_md)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Complete. {len(results)} file(s) processed.")
|
||||
print("Output files:")
|
||||
for r in results:
|
||||
print(f" {r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
353
scripts/step1_ocr_wardiaries_olmocr.py
Normal file
353
scripts/step1_ocr_wardiaries_olmocr.py
Normal file
@@ -0,0 +1,353 @@
|
||||
"""
|
||||
ocr_wardiaries.py
|
||||
-----------------
|
||||
Converts war diary PDFs to structured text using the olmOCR model hosted on DeepInfra.
|
||||
Bypasses the olmOCR pipeline GPU check by calling the API directly.
|
||||
|
||||
Output format: one .md file per PDF with ## Page N markers.
|
||||
- Diary pages with Place/Date/Hour/Summary tables are output as HTML tables,
|
||||
preserving column structure for downstream extraction.
|
||||
- All other pages (admin, appendices, forms, etc.) are output as plain text.
|
||||
|
||||
This structured output is the source of truth for the pipeline:
|
||||
- llm_extract.py reads the HTML tables for position and narrative extraction
|
||||
- json_to_viewer_md.py flattens the HTML tables to readable prose for the OCR viewer
|
||||
|
||||
Requirements:
|
||||
pip install requests pillow python-dotenv
|
||||
|
||||
Poppler must be on PATH (pdftoppm command must work).
|
||||
|
||||
API key is read from DEEPINFRA_API_KEY in your .env file.
|
||||
|
||||
Usage:
|
||||
# Test run — pages 7 to 27 only:
|
||||
python ocr_wardiaries.py --single "G:/path/to/diary.pdf" --output_dir "G:/path/to/outputs" --first_page 7 --last_page 27
|
||||
|
||||
# Full diary:
|
||||
python ocr_wardiaries.py --single "G:/path/to/diary.pdf" --output_dir "G:/path/to/outputs"
|
||||
|
||||
# All PDFs in a folder:
|
||||
python ocr_wardiaries.py --input_dir "G:/path/to/pdfs" --output_dir "G:/path/to/outputs"
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
import requests
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||||
MODEL = "allenai/olmOCR-2-7B-1025"
|
||||
DPI = 150
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 15
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OCR prompt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are transcribing a scanned WWII Canadian Army war diary page.
|
||||
|
||||
Transcribe every word exactly as written. Do not correct spelling, grammar,
|
||||
or abbreviations. Do not add commentary or explanations.
|
||||
Keep grid references exactly as written (e.g. MR 2553, GR 442891, MR 0675 Sheet 861).
|
||||
Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O., RHC, R de Mais).
|
||||
Transcribe all text visible on the page — headers, titles, form numbers,
|
||||
stamps, handwritten notes, signatures, and margin notations. Nothing omitted.
|
||||
|
||||
If the page contains a table with columns Place / Date / Hour / Summary / Remarks:
|
||||
Output the table as HTML using <table>, <tr>, <th>, <td> tags.
|
||||
- Place column: all lines exactly as written, separated by <br> tags.
|
||||
Include country names, place names, grid references, and sheet references.
|
||||
- Date column: exact date as written, <br> if split across lines.
|
||||
- Hour column: exact time as written, or empty if absent.
|
||||
- Summary column: full verbatim text, paragraphs separated by <br>. Every word. Nothing omitted.
|
||||
- Remarks column: exact text, or empty if absent.
|
||||
If the page has multiple rows (different dates or continuation entries), each row is a separate <tr>.
|
||||
Any text outside the table (headers, weather notes, stamps) is output as plain text
|
||||
before or after the <table> tag in the order it appears on the page.
|
||||
|
||||
If the page does NOT contain a Place/Date/Hour/Summary table:
|
||||
Output all text as plain text in the order it appears on the page.
|
||||
|
||||
If the page is genuinely blank, output only: [BLANK]\
|
||||
"""
|
||||
|
||||
USER_PROMPT = (
|
||||
"Transcribe every word on this page exactly as written. "
|
||||
"Use an HTML table for any Place/Date/Hour/Summary diary table. "
|
||||
"Nothing omitted, nothing added."
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI,
|
||||
first_page: int = None, last_page: int = None) -> list[tuple[int, Path]]:
|
||||
"""
|
||||
Render pages of a PDF to PNG images using pdftoppm (poppler).
|
||||
Returns a list of (pdf_page_number, image_path) tuples.
|
||||
"""
|
||||
prefix = output_dir / pdf_path.stem
|
||||
cmd = ["pdftoppm", "-r", str(dpi), "-png"]
|
||||
if first_page is not None:
|
||||
cmd += ["-f", str(first_page)]
|
||||
if last_page is not None:
|
||||
cmd += ["-l", str(last_page)]
|
||||
cmd += [str(pdf_path), str(prefix)]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"pdftoppm failed for {pdf_path.name}:\n{result.stderr}")
|
||||
|
||||
images = sorted(output_dir.glob(f"{pdf_path.stem}-*.png"))
|
||||
result_pairs = []
|
||||
for img_path in images:
|
||||
m = re.search(r'-(\d+)\.png$', img_path.name)
|
||||
page_num = int(m.group(1)) if m else len(result_pairs) + (first_page or 1)
|
||||
result_pairs.append((page_num, img_path))
|
||||
return result_pairs
|
||||
|
||||
|
||||
def image_to_base64(image_path: Path) -> str:
|
||||
"""Convert an image file to a base64 string."""
|
||||
with open(image_path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
|
||||
"""Send one page image to DeepInfra olmOCR and return the text string."""
|
||||
b64 = image_to_base64(image_path)
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": SYSTEM_PROMPT,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{b64}"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": USER_PROMPT,
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
DEEPINFRA_API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=180,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data["choices"][0]["message"]["content"].strip()
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — HTTP error: {e}]"
|
||||
except Exception as e:
|
||||
print(f" Error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED — page {page_num} — Error: {e}]"
|
||||
|
||||
|
||||
def has_table(page_text: str) -> bool:
|
||||
"""Return True if the page output contains an HTML table."""
|
||||
return bool(re.search(r'<table', page_text, re.IGNORECASE))
|
||||
|
||||
|
||||
def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str,
|
||||
first_page: int = None, last_page: int = None) -> Path:
|
||||
"""
|
||||
OCR a single PDF and write structured output to a .md file.
|
||||
Diary pages are output as HTML tables; all other pages as plain text.
|
||||
"""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Processing: {pdf_path.name}")
|
||||
if first_page or last_page:
|
||||
print(f"Pages: {first_page or 1} → {last_page or 'end'}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
output_md = output_dir / f"step1_{pdf_path.stem}_olmocr.md"
|
||||
|
||||
if output_md.exists():
|
||||
print(f" Already processed — skipping.")
|
||||
print(f" Delete {output_md.name} to reprocess.")
|
||||
return output_md
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
print(f" Rendering PDF pages to images at {DPI} DPI...")
|
||||
try:
|
||||
page_images = pdf_to_images(pdf_path, tmp_path,
|
||||
first_page=first_page, last_page=last_page)
|
||||
except RuntimeError as e:
|
||||
print(f" ERROR: {e}")
|
||||
return output_md
|
||||
|
||||
total = len(page_images)
|
||||
print(f" Found {total} pages to process.")
|
||||
|
||||
lines = [
|
||||
f"# {pdf_path.stem}",
|
||||
f"OCR by olmOCR ({MODEL}) via DeepInfra",
|
||||
f"Source: {pdf_path.name} — pages {first_page or 1}–{last_page or total}",
|
||||
"",
|
||||
]
|
||||
|
||||
diary_count = 0
|
||||
other_count = 0
|
||||
|
||||
for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
|
||||
print(f" Page {pdf_page_num} ({i}/{total})...", end=" ", flush=True)
|
||||
|
||||
page_text = ocr_page(image_path, api_key, pdf_page_num)
|
||||
is_diary = has_table(page_text)
|
||||
|
||||
lines.append(f"## Page {pdf_page_num}")
|
||||
lines.append("")
|
||||
lines.append(page_text)
|
||||
lines.append("")
|
||||
|
||||
if is_diary:
|
||||
diary_count += 1
|
||||
print("done — diary (table)")
|
||||
else:
|
||||
other_count += 1
|
||||
print("done — plain text")
|
||||
|
||||
if i < total:
|
||||
time.sleep(0.5)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_md.write_text("\n".join(lines), encoding="utf-8")
|
||||
|
||||
print(f"\n Saved: {output_md}")
|
||||
print(f" Summary: {diary_count} diary pages (HTML tables), {other_count} other pages (plain text)")
|
||||
|
||||
return output_md
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="OCR war diary PDFs to structured text using olmOCR via DeepInfra."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_dir",
|
||||
default=None,
|
||||
help="Folder containing PDF files to process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
required=True,
|
||||
help="Folder to write output .md files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--single",
|
||||
default=None,
|
||||
help="Process a single PDF file instead of a whole folder.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--first_page",
|
||||
type=int,
|
||||
default=None,
|
||||
help="First PDF page to process (1-based).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--last_page",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Last PDF page to process (1-based).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = os.getenv("DEEPINFRA_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: DEEPINFRA_API_KEY not set in environment or .env file")
|
||||
sys.exit(1)
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if args.single:
|
||||
pdfs = [Path(args.single)]
|
||||
elif args.input_dir:
|
||||
pdfs = sorted(Path(args.input_dir).glob("*.pdf"))
|
||||
else:
|
||||
print("Error: provide --input_dir or --single")
|
||||
sys.exit(1)
|
||||
|
||||
if not pdfs:
|
||||
print("No PDF files found.")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Found {len(pdfs)} PDF(s) to process.")
|
||||
print(f"Output directory: {output_dir}")
|
||||
print(f"Model: {MODEL}")
|
||||
print(f"DPI: {DPI}")
|
||||
|
||||
results = []
|
||||
for pdf_path in pdfs:
|
||||
output_md = ocr_pdf(
|
||||
pdf_path, output_dir, api_key,
|
||||
first_page=args.first_page,
|
||||
last_page=args.last_page,
|
||||
)
|
||||
results.append(output_md)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Complete. {len(results)} file(s) processed.")
|
||||
print("Output files:")
|
||||
for r in results:
|
||||
print(f" {r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
327
scripts/step2_json_to_viewer_md.py
Normal file
327
scripts/step2_json_to_viewer_md.py
Normal file
@@ -0,0 +1,327 @@
|
||||
"""
|
||||
json_to_viewer_md.py
|
||||
--------------------
|
||||
Converts the structured olmOCR output (.md with HTML tables) into clean
|
||||
readable prose for the OCR viewer.
|
||||
|
||||
Reads the _olmocr.md file produced by ocr_wardiaries.py.
|
||||
- Diary pages (HTML tables) are flattened to readable prose.
|
||||
- Non-diary pages (plain text) are passed through unchanged.
|
||||
|
||||
Output: a clean .md file with ## Page N markers, no HTML tags,
|
||||
suitable for display in the OCR viewer.
|
||||
|
||||
No API calls. Pure Python. Runs locally in seconds.
|
||||
|
||||
Usage:
|
||||
python json_to_viewer_md.py --input "G:/path/to/outputs_step1_olmocr/Calgary-Highlanders_War-Diary_Sep44_olmocr.md" --output_dir "G:/path/to/outputs_step2_json_to_viewer_md"
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTML table parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TableParser(HTMLParser):
|
||||
"""
|
||||
Parses a single HTML table into a list of row dicts.
|
||||
Each row: { place, date, hour, summary, remarks }
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.rows = []
|
||||
self._in_table = False
|
||||
self._in_header = False
|
||||
self._current_row = None
|
||||
self._current_cell = None
|
||||
self._cell_index = 0
|
||||
self._cell_parts = []
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
tag = tag.lower()
|
||||
if tag == 'table':
|
||||
self._in_table = True
|
||||
elif tag == 'tr' and self._in_table:
|
||||
self._current_row = []
|
||||
self._cell_index = 0
|
||||
elif tag in ('td', 'th') and self._current_row is not None:
|
||||
self._in_header = (tag == 'th')
|
||||
self._current_cell = []
|
||||
self._cell_parts = []
|
||||
elif tag == 'br' and self._current_cell is not None:
|
||||
self._cell_parts.append('\n')
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
tag = tag.lower()
|
||||
if tag in ('td', 'th') and self._current_cell is not None:
|
||||
cell_text = ''.join(self._cell_parts).strip()
|
||||
self._current_row.append(cell_text)
|
||||
self._current_cell = None
|
||||
self._cell_parts = []
|
||||
self._cell_index += 1
|
||||
elif tag == 'tr' and self._current_row is not None:
|
||||
if not self._in_header and len(self._current_row) >= 4:
|
||||
self.rows.append({
|
||||
'place': self._current_row[0] if len(self._current_row) > 0 else '',
|
||||
'date': self._current_row[1] if len(self._current_row) > 1 else '',
|
||||
'hour': self._current_row[2] if len(self._current_row) > 2 else '',
|
||||
'summary': self._current_row[3] if len(self._current_row) > 3 else '',
|
||||
'remarks': self._current_row[4] if len(self._current_row) > 4 else '',
|
||||
})
|
||||
self._current_row = None
|
||||
self._in_header = False
|
||||
elif tag == 'table':
|
||||
self._in_table = False
|
||||
|
||||
def handle_data(self, data):
|
||||
if self._current_cell is not None:
|
||||
self._cell_parts.append(data)
|
||||
|
||||
|
||||
def parse_table(html: str) -> list[dict]:
|
||||
"""Parse an HTML table string into a list of row dicts."""
|
||||
parser = TableParser()
|
||||
parser.feed(html)
|
||||
return parser.rows
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Text cleaning
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def clean_cell(text: str) -> str:
|
||||
"""Strip leading/trailing whitespace and collapse multiple newlines to one."""
|
||||
text = re.sub(r'\n{2,}', '\n', text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def format_place(place: str) -> str:
|
||||
"""
|
||||
Format the Place column for readability.
|
||||
Lines separated by newlines → joined with ' / ' for display.
|
||||
Removes blank lines.
|
||||
"""
|
||||
lines = [l.strip() for l in place.split('\n') if l.strip()]
|
||||
return ' / '.join(lines)
|
||||
|
||||
|
||||
def format_date(date: str) -> str:
|
||||
"""Collapse multiline date (e.g. '1 Sep\n44') to single line."""
|
||||
parts = [p.strip() for p in date.split('\n') if p.strip()]
|
||||
return ' '.join(parts)
|
||||
|
||||
|
||||
def row_to_prose(row: dict) -> str:
|
||||
"""
|
||||
Convert a single diary table row to readable prose for the viewer.
|
||||
|
||||
Format:
|
||||
DATE — PLACE
|
||||
[HOUR]
|
||||
SUMMARY TEXT
|
||||
|
||||
Remarks: REMARKS
|
||||
"""
|
||||
parts = []
|
||||
|
||||
date = format_date(clean_cell(row.get('date', '')))
|
||||
place = format_place(clean_cell(row.get('place', '')))
|
||||
hour = clean_cell(row.get('hour', ''))
|
||||
summary = clean_cell(row.get('summary', ''))
|
||||
remarks = clean_cell(row.get('remarks', ''))
|
||||
|
||||
# Header line: date and place
|
||||
header_parts = []
|
||||
if date:
|
||||
header_parts.append(date)
|
||||
if place:
|
||||
header_parts.append(place)
|
||||
if header_parts:
|
||||
parts.append(' — '.join(header_parts))
|
||||
|
||||
# Hour on its own line if present
|
||||
if hour:
|
||||
parts.append(f'Hour: {hour}')
|
||||
|
||||
# Blank line before summary
|
||||
if summary:
|
||||
parts.append('')
|
||||
parts.append(summary)
|
||||
|
||||
# Remarks if present
|
||||
if remarks:
|
||||
parts.append('')
|
||||
parts.append(f'Remarks: {remarks}')
|
||||
|
||||
return '\n'.join(parts)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MD page parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_md_pages(md_text: str) -> dict[int, str]:
|
||||
"""
|
||||
Parse a .md file into a dict of {page_number: page_text}.
|
||||
"""
|
||||
pages = {}
|
||||
page_re = re.compile(r'^## Page (\d+)\s*$', re.MULTILINE)
|
||||
matches = list(page_re.finditer(md_text))
|
||||
|
||||
for i, m in enumerate(matches):
|
||||
page_num = int(m.group(1))
|
||||
start = m.end()
|
||||
end = matches[i + 1].start() if i + 1 < len(matches) else len(md_text)
|
||||
page_text = md_text[start:end].strip()
|
||||
pages[page_num] = page_text
|
||||
|
||||
return pages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Page flattener
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def flatten_page(page_num: int, page_text: str) -> str:
|
||||
"""
|
||||
Flatten one page of OCR output to clean prose for the viewer.
|
||||
|
||||
- If the page contains HTML tables: parse and flatten each table row.
|
||||
- If the page is plain text: pass through as-is.
|
||||
"""
|
||||
|
||||
# Extract any weather line that sits outside the table
|
||||
weather_match = re.search(
|
||||
r'(?:^|\n)(Weather[:\s].+?)(?:\n|$)',
|
||||
page_text,
|
||||
re.IGNORECASE
|
||||
)
|
||||
weather_line = weather_match.group(1).strip() if weather_match else None
|
||||
|
||||
# Find all tables on this page
|
||||
table_re = re.compile(r'<table[\s\S]*?</table>', re.IGNORECASE)
|
||||
tables = table_re.findall(page_text)
|
||||
|
||||
if not tables:
|
||||
# No table — plain text page, pass through unchanged
|
||||
return page_text
|
||||
|
||||
# Diary page — parse each table and flatten rows
|
||||
output_parts = []
|
||||
|
||||
for table_html in tables:
|
||||
rows = parse_table(table_html)
|
||||
if not rows:
|
||||
continue
|
||||
for row in rows:
|
||||
prose = row_to_prose(row)
|
||||
if prose.strip():
|
||||
output_parts.append(prose)
|
||||
output_parts.append('') # blank line between entries
|
||||
|
||||
# Append weather if present
|
||||
if weather_line:
|
||||
output_parts.append(weather_line)
|
||||
|
||||
return '\n'.join(output_parts).strip()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main converter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def convert(input_path: Path, output_dir: Path):
|
||||
"""
|
||||
Read the olmOCR .md file and write a clean viewer .md file.
|
||||
"""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Converting: {input_path.name}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
md_text = input_path.read_text(encoding='utf-8')
|
||||
pages = parse_md_pages(md_text)
|
||||
|
||||
if not pages:
|
||||
print(" ERROR: No pages found. Check file format.")
|
||||
sys.exit(1)
|
||||
|
||||
print(f" Found {len(pages)} pages.")
|
||||
|
||||
# Build output
|
||||
stem = input_path.stem
|
||||
output_path = output_dir / f"{stem}_viewer.md"
|
||||
|
||||
lines = [
|
||||
f"# {stem}",
|
||||
f"Generated by json_to_viewer_md.py from {input_path.name}",
|
||||
"",
|
||||
]
|
||||
|
||||
diary_count = 0
|
||||
plain_count = 0
|
||||
|
||||
for page_num in sorted(pages.keys()):
|
||||
page_text = pages[page_num]
|
||||
|
||||
lines.append(f"## Page {page_num}")
|
||||
lines.append("")
|
||||
|
||||
if not page_text.strip() or page_text.strip() == '[BLANK]':
|
||||
lines.append('[BLANK]')
|
||||
plain_count += 1
|
||||
else:
|
||||
flattened = flatten_page(page_num, page_text)
|
||||
lines.append(flattened)
|
||||
if re.search(r'<table', page_text, re.IGNORECASE):
|
||||
diary_count += 1
|
||||
else:
|
||||
plain_count += 1
|
||||
|
||||
lines.append("")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text('\n'.join(lines), encoding='utf-8')
|
||||
|
||||
print(f" Diary pages flattened : {diary_count}")
|
||||
print(f" Plain text pages : {plain_count}")
|
||||
print(f" Saved: {output_path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Flatten olmOCR HTML table output to clean prose for the OCR viewer."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
required=True,
|
||||
help="Path to the _olmocr.md file from ocr_wardiaries.py.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
required=True,
|
||||
help="Folder to write the _viewer.md file.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
input_path = Path(args.input)
|
||||
if not input_path.exists():
|
||||
print(f"Error: input file not found: {input_path}")
|
||||
sys.exit(1)
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
convert(input_path, output_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
600
scripts/step3_extract-posn_claude-4-6.py
Normal file
600
scripts/step3_extract-posn_claude-4-6.py
Normal file
@@ -0,0 +1,600 @@
|
||||
"""
|
||||
step3_extract-posn_claude-4-6.py
|
||||
----------------------------------
|
||||
Extracts structured data from verified OCR text produced by
|
||||
step1_ocr_wardiaries_claude-4-6.py (or the olmOCR equivalent).
|
||||
|
||||
Handles two page types:
|
||||
|
||||
DIARY PAGES (HTML tables — Place / Date / Hour / Summary / Remarks):
|
||||
Extracts dated narrative entries and all position records.
|
||||
Output: {stem}_claude-4-6_narratives.json
|
||||
{stem}_claude-4-6_positions.json
|
||||
|
||||
APPENDIX PAGES (plain text — message forms, movement orders, arty tables,
|
||||
field returns, ISUMs, patrol reports, traces, etc.):
|
||||
Extracts document metadata and all grid references / positions mentioned.
|
||||
Output: {stem}_claude-4-6_appx_records.json
|
||||
{stem}_claude-4-6_appx_positions.json
|
||||
|
||||
Requirements:
|
||||
pip install requests python-dotenv
|
||||
|
||||
API key is read from ANTHROPIC_API_KEY in your .env file.
|
||||
|
||||
Usage:
|
||||
python scripts/step3_extract-posn_claude-4-6.py --input "G:/path/to/step1_Calgary-Highlanders_War-Diary_Sep44_claude-4-6.md" --output_dir "G:/path/to/outputs_step3_llm"
|
||||
python scripts/step3_extract-posn_claude-4-6.py --input "..." --output_dir "..." --first_page 7 --last_page 27
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
import requests
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
|
||||
MODEL = "claude-sonnet-4-6"
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 10
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are a military historian extracting structured data from a WWII Canadian Army
|
||||
war diary file (Calgary Highlanders, September–October 1944).
|
||||
|
||||
You will receive the raw text of one page. Determine the page type and extract
|
||||
accordingly. Output valid JSON only. No prose before or after the JSON.
|
||||
|
||||
════════════════════════════════════════
|
||||
PAGE TYPE DETECTION
|
||||
════════════════════════════════════════
|
||||
|
||||
page_type must be one of:
|
||||
"diary" — HTML table with Place / Date / Hour / Summary / Remarks columns
|
||||
"message_form" — Army Form C2136 or similar signal/message form
|
||||
"movement_order" — Numbered movement or operation order
|
||||
"arty_table" — Artillery DF / SOS / fire task table
|
||||
"field_return" — Strength return (officers or other ranks)
|
||||
"isum" — Intelligence summary
|
||||
"patrol_report" — Patrol programme or patrol report
|
||||
"trace" — Map, sketch map, disposition diagram, or trace
|
||||
"admin" — Part I/II orders, nominal rolls, boilerplate instructions
|
||||
"other" — Anything that does not fit the above
|
||||
|
||||
════════════════════════════════════════
|
||||
OUTPUT SCHEMA — ALL PAGES
|
||||
════════════════════════════════════════
|
||||
|
||||
{
|
||||
"page_type": "diary | message_form | movement_order | ...",
|
||||
"date_warning": null or "explanation",
|
||||
|
||||
// ── DIARY PAGES ONLY ──────────────────────────────────────────────────────
|
||||
"narratives": [
|
||||
{
|
||||
"date": "1 Sep 44",
|
||||
"date_inferred": false,
|
||||
"date_span": null,
|
||||
"place": "Ste. Foy",
|
||||
"summary": "Full verbatim narrative text...",
|
||||
"weather": "Fair and warm",
|
||||
"no_change": false,
|
||||
"source_page": 7
|
||||
}
|
||||
],
|
||||
|
||||
// ── APPENDIX PAGES ONLY ───────────────────────────────────────────────────
|
||||
"appx_record": {
|
||||
"page_type": "message_form",
|
||||
"document_id": "GO-7",
|
||||
"date": "27 Sep 44",
|
||||
"time": "2205A",
|
||||
"from_unit": "HQ RCA 2 Cdn Inf Div",
|
||||
"to_units": ["5 CIB", "RHC", "CALG HIGHRS"],
|
||||
"subject": "DF Task Table No 35 — amendment",
|
||||
"summary": "One sentence description of what this document contains.",
|
||||
"units_mentioned": ["Calgary Highlanders", "RHC", "R de Mais", "5 Cdn Fd Regt"],
|
||||
"source_page": 12
|
||||
},
|
||||
|
||||
// ── ALL PAGES — positions found anywhere on the page ─────────────────────
|
||||
"positions": [
|
||||
{
|
||||
"date": "1 Sep 44",
|
||||
"date_inferred": false,
|
||||
"hour": null,
|
||||
"grid": "255330",
|
||||
"grid_inferred": false,
|
||||
"place_name": "Ste. Foy",
|
||||
"sheet_ref": "Sheet 861",
|
||||
"category": "END_OF_DAY",
|
||||
"subunit": null,
|
||||
"friendly_unit": null,
|
||||
"no_movement": false,
|
||||
"confidence": "high",
|
||||
"context": "arrived in the little village of Ste. Foy east of Longueville",
|
||||
"source_page": 7
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
For diary pages: populate "narratives" and "positions". Set "appx_record" to null.
|
||||
For appendix pages: populate "appx_record" and "positions". Set "narratives" to [].
|
||||
For blank pages: set all arrays empty, appx_record null.
|
||||
|
||||
════════════════════════════════════════
|
||||
DIARY PAGE RULES (page_type = "diary")
|
||||
════════════════════════════════════════
|
||||
|
||||
Each <tr> in the HTML table is one diary entry. Extract:
|
||||
|
||||
narratives:
|
||||
date — from Date <td>. Join multiline ("1 Sep\n44" → "1 Sep 44").
|
||||
date_span — if entry covers multiple dates ("6-8 Sep"), set to "6-8 Sep 44",
|
||||
date to the first date.
|
||||
date_inferred — true if date carried forward from a previous row.
|
||||
place — primary place name from Place <td>, excluding grid/sheet refs.
|
||||
summary — full verbatim text from Summary <td>. Nothing omitted.
|
||||
weather — weather note if present on the page. null if absent.
|
||||
no_change — true if entry states no change from previous day.
|
||||
|
||||
positions (cast a wide net — extract from BOTH Place <td> AND Summary <td>):
|
||||
PLACE COLUMN → category END_OF_DAY (last entry per date) or UNIT_MOVEMENT.
|
||||
SUMMARY TEXT → category SUBUNIT, FRIENDLY, ENEMY, PATROL, or MISC.
|
||||
Only ONE position per date may have category END_OF_DAY.
|
||||
|
||||
════════════════════════════════════════
|
||||
APPENDIX PAGE RULES (all other page_types)
|
||||
════════════════════════════════════════
|
||||
|
||||
appx_record:
|
||||
document_id — order/form number if visible (e.g. "Mov Order No 5",
|
||||
"ISUM No 45", "GO-3", "DF Task Table No 35"). null if absent.
|
||||
date — date of the document. null if not determinable.
|
||||
time — time or date-time group if present (e.g. "272205A"). null if absent.
|
||||
from_unit — originating unit/HQ. null if absent.
|
||||
to_units — list of addressees. [] if absent.
|
||||
subject — subject line or a one-sentence description of purpose.
|
||||
summary — one to three sentence plain-English summary of what this
|
||||
document records or orders. No verbatim transcription.
|
||||
units_mentioned — all unit names appearing anywhere on the page.
|
||||
|
||||
positions (extract ALL grid references and named locations):
|
||||
Extract every grid reference (4-digit, 6-digit, 8-digit) and every named
|
||||
location from the entire page, regardless of context.
|
||||
Use the same position schema as diary pages.
|
||||
category:
|
||||
UNIT_MOVEMENT — if associated with the Calgary Highlanders' own movement
|
||||
FRIENDLY — if associated with another Allied unit
|
||||
ENEMY — if associated with enemy forces
|
||||
PATROL — if from a patrol programme or patrol report
|
||||
DF_TASK — if from an DF/SOS artillery task table
|
||||
MISC — everything else (objectives, named features, route points)
|
||||
date — take from document date if not stated per-position. null if unknown.
|
||||
confidence:
|
||||
"high" — grid explicitly written next to a place name
|
||||
"medium" — grid present but context unclear
|
||||
"low" — place name only, no grid; or grid with no place name context
|
||||
|
||||
════════════════════════════════════════
|
||||
GRID REFERENCE RULES (all page types)
|
||||
════════════════════════════════════════
|
||||
|
||||
grid: digits only, no prefix, no punctuation.
|
||||
"MR 2468" → "2468" "GR 442891" → "442891" "MR 24.68" → "2468"
|
||||
If the diary gives a 4-digit grid, expand: "2553" → "255535". Set grid_inferred = true.
|
||||
Do not guess or expand 6-digit grids. Set grid_inferred = false.
|
||||
8-digit grids: keep first 3 + last 3 digits → "44289100" → "442891".
|
||||
5-digit grids: set grid = null, confidence = "low", note in context.
|
||||
If no grid is available, set grid to null and record place_name instead.
|
||||
|
||||
sheet_ref: carry forward the most recent sheet reference seen on the page.
|
||||
hour: HHMM format if stated ("0930"). null if absent.
|
||||
Ignore Hour <td> values that look like years ("44") or grid refs.
|
||||
|
||||
Output ONLY the JSON object. No preamble. No explanation. No trailing text.\
|
||||
"""
|
||||
|
||||
|
||||
def make_user_prompt(page_num: int, page_text: str) -> str:
|
||||
return (
|
||||
f"This is page {page_num} of the war diary file. "
|
||||
f"Determine the page type, then extract accordingly.\n\n"
|
||||
f"PAGE TEXT:\n{page_text}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MD page parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_md_pages(md_text: str) -> dict[int, str]:
|
||||
pages = {}
|
||||
page_re = re.compile(r'^## Page (\d+)\s*$', re.MULTILINE)
|
||||
matches = list(page_re.finditer(md_text))
|
||||
for i, m in enumerate(matches):
|
||||
page_num = int(m.group(1))
|
||||
start = m.end()
|
||||
end = matches[i + 1].start() if i + 1 < len(matches) else len(md_text)
|
||||
pages[page_num] = md_text[start:end].strip()
|
||||
return pages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON response parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_json_response(raw: str, page_num: int) -> dict | None:
|
||||
cleaned = raw.strip()
|
||||
cleaned = re.sub(r'^```json\s*', '', cleaned, flags=re.IGNORECASE)
|
||||
cleaned = re.sub(r'^```\s*', '', cleaned)
|
||||
cleaned = re.sub(r'\s*```$', '', cleaned)
|
||||
cleaned = cleaned.strip()
|
||||
try:
|
||||
return json.loads(cleaned)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f" WARNING: JSON parse failed for page {page_num}: {e}")
|
||||
print(f" Raw response (first 300 chars): {raw[:300]}")
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API call — Anthropic
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def extract_page(page_num: int, page_text: str, api_key: str) -> dict | None:
|
||||
"""Send one page to Claude via Anthropic API and return parsed extraction result."""
|
||||
|
||||
headers = {
|
||||
"x-api-key": api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"content-type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"max_tokens": 8192,
|
||||
"system": SYSTEM_PROMPT,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": make_user_prompt(page_num, page_text),
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
ANTHROPIC_API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
usage = data.get("usage", {})
|
||||
input_t = usage.get("input_tokens", 0)
|
||||
output_t = usage.get("output_tokens", 0)
|
||||
cost = (input_t / 1_000_000 * 3.00) + (output_t / 1_000_000 * 15.00)
|
||||
print(f" tokens: {input_t} in / {output_t} out — ${cost:.4f}", end=" ")
|
||||
|
||||
raw_content = data["content"][0]["text"]
|
||||
return parse_json_response(raw_content, page_num)
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if hasattr(e, 'response') and e.response is not None:
|
||||
print(f" Response: {e.response.text[:300]}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" Error on page {page_num, attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-of-day assignment (diary positions only)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def assign_end_of_day(positions: list[dict]) -> list[dict]:
|
||||
from collections import defaultdict
|
||||
date_groups: dict[str, list[int]] = defaultdict(list)
|
||||
for i, p in enumerate(positions):
|
||||
if p.get('date'):
|
||||
date_groups[p['date']].append(i)
|
||||
|
||||
for date, indices in date_groups.items():
|
||||
eod_indices = [i for i in indices if positions[i].get('category') == 'END_OF_DAY']
|
||||
if len(eod_indices) == 1:
|
||||
continue
|
||||
if len(eod_indices) > 1:
|
||||
for i in eod_indices[:-1]:
|
||||
positions[i]['category'] = 'UNIT_MOVEMENT'
|
||||
continue
|
||||
with_grid = [i for i in indices if positions[i].get('grid')]
|
||||
unit_movement = [i for i in with_grid if positions[i].get('category') == 'UNIT_MOVEMENT']
|
||||
if unit_movement:
|
||||
positions[unit_movement[-1]]['category'] = 'END_OF_DAY'
|
||||
elif with_grid:
|
||||
positions[with_grid[-1]]['category'] = 'END_OF_DAY'
|
||||
return positions
|
||||
|
||||
|
||||
def post_process_positions(positions: list[dict]) -> list[dict]:
|
||||
for p in positions:
|
||||
raw = re.sub(r'[^0-9]', '', p.get('grid') or '')
|
||||
if len(raw) == 4:
|
||||
p['grid'] = raw[0:2] + '5' + raw[2:4] + '5'
|
||||
p['grid_inferred'] = True
|
||||
elif len(raw) == 6:
|
||||
p['grid'] = raw
|
||||
p['grid_inferred'] = False
|
||||
elif len(raw) == 8:
|
||||
p['grid'] = raw[0:3] + raw[4:7]
|
||||
p['grid_inferred'] = False
|
||||
elif len(raw) == 5:
|
||||
p['grid'] = None
|
||||
p['grid_inferred'] = False
|
||||
p['confidence'] = 'low'
|
||||
p['context'] = (p.get('context') or '') + ' [5-figure grid — needs human review]'
|
||||
else:
|
||||
p['grid'] = None
|
||||
p['grid_inferred'] = False
|
||||
p['confidence'] = 'low'
|
||||
return positions
|
||||
|
||||
|
||||
def dedup_positions(positions: list[dict]) -> list[dict]:
|
||||
seen = set()
|
||||
result = []
|
||||
for p in positions:
|
||||
key = (p.get('date'), p.get('grid'), (p.get('context') or '')[:60])
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
result.append(p)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main extraction loop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def extract_diary(md_path: Path, output_dir: Path, api_key: str,
|
||||
first_page: int = None, last_page: int = None):
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Extracting: {md_path.name}")
|
||||
if first_page or last_page:
|
||||
print(f"Pages: {first_page or 'start'} → {last_page or 'end'}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
stem = md_path.stem
|
||||
for suffix in ('_olmocr', '_claude-4-6', 'step1_'):
|
||||
stem = stem.replace(suffix, '')
|
||||
stem = stem.strip('_')
|
||||
|
||||
# Diary outputs
|
||||
narratives_out = output_dir / f"{stem}_claude-4-6_narratives.json"
|
||||
positions_out = output_dir / f"{stem}_claude-4-6_positions.json"
|
||||
# Appendix outputs
|
||||
appx_records_out = output_dir / f"{stem}_claude-4-6_appx_records.json"
|
||||
appx_posns_out = output_dir / f"{stem}_claude-4-6_appx_positions.json"
|
||||
# Checkpoint file — deleted on clean completion
|
||||
checkpoint_path = output_dir / f"{stem}_claude-4-6_checkpoint.json"
|
||||
|
||||
md_text = md_path.read_text(encoding='utf-8')
|
||||
all_pages = parse_md_pages(md_text)
|
||||
|
||||
if not all_pages:
|
||||
print(" ERROR: No pages found in MD file. Check file format.")
|
||||
return
|
||||
|
||||
page_nums = sorted(all_pages.keys())
|
||||
if first_page:
|
||||
page_nums = [p for p in page_nums if p >= first_page]
|
||||
if last_page:
|
||||
page_nums = [p for p in page_nums if p <= last_page]
|
||||
|
||||
# ── Resume from checkpoint if one exists ─────────────────────────────────
|
||||
all_narratives = []
|
||||
all_positions = []
|
||||
all_appx_records = []
|
||||
all_appx_posns = []
|
||||
skipped_pages = []
|
||||
error_pages = []
|
||||
resume_from = None
|
||||
|
||||
if checkpoint_path.exists():
|
||||
try:
|
||||
ckpt = json.loads(checkpoint_path.read_text(encoding='utf-8'))
|
||||
resume_from = ckpt.get('last_completed_page')
|
||||
all_narratives = ckpt.get('narratives', [])
|
||||
all_positions = ckpt.get('positions', [])
|
||||
all_appx_records = ckpt.get('appx_records', [])
|
||||
all_appx_posns = ckpt.get('appx_positions',[])
|
||||
skipped_pages = ckpt.get('skipped_pages', [])
|
||||
error_pages = ckpt.get('error_pages', [])
|
||||
print(f" RESUMING from checkpoint — last completed page: {resume_from}")
|
||||
page_nums = [p for p in page_nums if p > resume_from]
|
||||
print(f" Remaining pages: {len(page_nums)}")
|
||||
except Exception as e:
|
||||
print(f" WARNING: Could not read checkpoint ({e}) — starting from scratch.")
|
||||
|
||||
print(f" Found {len(all_pages)} total pages, processing {len(page_nums)}.")
|
||||
print(f" Model: {MODEL}")
|
||||
print(f" Diary outputs : {narratives_out.name} / {positions_out.name}")
|
||||
print(f" Appendix outputs: {appx_records_out.name} / {appx_posns_out.name}")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def save_checkpoint(last_page_done: int):
|
||||
checkpoint_path.write_text(
|
||||
json.dumps({
|
||||
'last_completed_page': last_page_done,
|
||||
'narratives': all_narratives,
|
||||
'positions': all_positions,
|
||||
'appx_records': all_appx_records,
|
||||
'appx_positions':all_appx_posns,
|
||||
'skipped_pages': skipped_pages,
|
||||
'error_pages': error_pages,
|
||||
}, ensure_ascii=False),
|
||||
encoding='utf-8'
|
||||
)
|
||||
|
||||
for i, page_num in enumerate(page_nums, start=1):
|
||||
page_text = all_pages[page_num]
|
||||
|
||||
if not page_text.strip() or page_text.strip() == '[BLANK]':
|
||||
print(f" Page {page_num} ({i}/{len(page_nums)})... skipped (blank)")
|
||||
skipped_pages.append(page_num)
|
||||
save_checkpoint(page_num)
|
||||
continue
|
||||
|
||||
print(f" Page {page_num} ({i}/{len(page_nums)})...", end=" ", flush=True)
|
||||
|
||||
result = extract_page(page_num, page_text, api_key)
|
||||
|
||||
if result is None:
|
||||
print("ERROR — extraction failed")
|
||||
error_pages.append(page_num)
|
||||
save_checkpoint(page_num)
|
||||
continue
|
||||
|
||||
page_type = result.get('page_type', 'other')
|
||||
date_warning = result.get('date_warning')
|
||||
narratives = result.get('narratives', []) or []
|
||||
positions = result.get('positions', []) or []
|
||||
appx_record = result.get('appx_record')
|
||||
|
||||
for n in narratives:
|
||||
n['source_page'] = page_num
|
||||
for p in positions:
|
||||
p['source_page'] = page_num
|
||||
if appx_record:
|
||||
appx_record['source_page'] = page_num
|
||||
|
||||
if page_type == 'diary':
|
||||
all_narratives.extend(narratives)
|
||||
all_positions.extend(positions)
|
||||
parts = [f"diary — {len(narratives)} narrative(s), {len(positions)} position(s)"]
|
||||
else:
|
||||
if appx_record:
|
||||
all_appx_records.append(appx_record)
|
||||
all_appx_posns.extend(positions)
|
||||
doc_id = (appx_record or {}).get('document_id') or ''
|
||||
parts = [f"{page_type}{' — ' + doc_id if doc_id else ''} — {len(positions)} position(s)"]
|
||||
|
||||
if date_warning:
|
||||
parts.append(f"DATE WARNING: {date_warning}")
|
||||
print(", ".join(parts))
|
||||
|
||||
# Save checkpoint after every successfully processed page
|
||||
save_checkpoint(page_num)
|
||||
|
||||
if i < len(page_nums):
|
||||
time.sleep(0.5)
|
||||
|
||||
# Post-process diary positions
|
||||
print(f"\n Running end-of-day assignment (diary)...")
|
||||
all_positions = assign_end_of_day(all_positions)
|
||||
all_positions = post_process_positions(all_positions)
|
||||
all_positions = dedup_positions(all_positions)
|
||||
|
||||
# Post-process appendix positions
|
||||
all_appx_posns = post_process_positions(all_appx_posns)
|
||||
all_appx_posns = dedup_positions(all_appx_posns)
|
||||
|
||||
# Write final outputs
|
||||
narratives_out.write_text(
|
||||
json.dumps(all_narratives, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
positions_out.write_text(
|
||||
json.dumps(all_positions, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
appx_records_out.write_text(
|
||||
json.dumps(all_appx_records, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
appx_posns_out.write_text(
|
||||
json.dumps(all_appx_posns, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
|
||||
# Clean up checkpoint — run completed successfully
|
||||
if checkpoint_path.exists():
|
||||
checkpoint_path.unlink()
|
||||
print(f" Checkpoint removed.")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Complete.")
|
||||
print(f" Diary narratives : {len(all_narratives)} entries → {narratives_out.name}")
|
||||
print(f" Diary positions : {len(all_positions)} records → {positions_out.name}")
|
||||
print(f" Appendix records : {len(all_appx_records)} documents → {appx_records_out.name}")
|
||||
print(f" Appendix positions: {len(all_appx_posns)} records → {appx_posns_out.name}")
|
||||
if skipped_pages:
|
||||
print(f" Skipped : pages {skipped_pages} (blank)")
|
||||
if error_pages:
|
||||
print(f" ERRORS : pages {error_pages} — re-run with --first_page/--last_page to retry")
|
||||
|
||||
warnings = [(n.get('source_page'), n.get('date'), n.get('date_warning'))
|
||||
for n in all_narratives if n.get('date_warning')]
|
||||
if warnings:
|
||||
print(f"\n DATE WARNINGS ({len(warnings)}):")
|
||||
for page, date, warning in warnings:
|
||||
print(f" Page {page} | date={date} | {warning}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Extract narratives, positions, and appendix records from OCR .md files using Claude Sonnet 4.6."
|
||||
)
|
||||
parser.add_argument("--input", required=True,
|
||||
help="Path to the step1_*_claude-4-6.md (or _olmocr.md) file.")
|
||||
parser.add_argument("--output_dir", required=True,
|
||||
help="Folder to write output JSON files.")
|
||||
parser.add_argument("--first_page", type=int, default=None,
|
||||
help="First page to process (1-based).")
|
||||
parser.add_argument("--last_page", type=int, default=None,
|
||||
help="Last page to process (1-based).")
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: ANTHROPIC_API_KEY not set in environment or .env file")
|
||||
sys.exit(1)
|
||||
|
||||
md_path = Path(args.input)
|
||||
if not md_path.exists():
|
||||
print(f"Error: input file not found: {md_path}")
|
||||
sys.exit(1)
|
||||
|
||||
extract_diary(
|
||||
md_path, Path(args.output_dir), api_key,
|
||||
first_page=args.first_page,
|
||||
last_page=args.last_page,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
599
scripts/step3_extract-posn_llama-3-1-70b-instruct.py
Normal file
599
scripts/step3_extract-posn_llama-3-1-70b-instruct.py
Normal file
@@ -0,0 +1,599 @@
|
||||
"""
|
||||
step3_extract-posn_llama-3-1-70b-instruct.py
|
||||
---------------------------------------------
|
||||
Extracts structured data from verified OCR text produced by
|
||||
step1_ocr_wardiaries_olmocr.py (or the claude-4-6 equivalent).
|
||||
|
||||
Handles two page types:
|
||||
|
||||
DIARY PAGES (HTML tables — Place / Date / Hour / Summary / Remarks):
|
||||
Extracts dated narrative entries and all position records.
|
||||
Output: {stem}_llama-3-1-70b_narratives.json
|
||||
{stem}_llama-3-1-70b_positions.json
|
||||
|
||||
APPENDIX PAGES (plain text — message forms, movement orders, arty tables,
|
||||
field returns, ISUMs, patrol reports, traces, etc.):
|
||||
Extracts document metadata and all grid references / positions mentioned.
|
||||
Output: {stem}_llama-3-1-70b_appx_records.json
|
||||
{stem}_llama-3-1-70b_appx_positions.json
|
||||
|
||||
Requirements:
|
||||
pip install requests python-dotenv
|
||||
|
||||
API key is read from DEEPINFRA_API_KEY in your .env file.
|
||||
|
||||
Usage:
|
||||
python scripts/step3_extract-posn_llama-3-1-70b-instruct.py --input "G:/path/to/step1_Calgary-Highlanders_War-Diary_Sep44_olmocr.md" --output_dir "G:/path/to/outputs_step3_llm"
|
||||
python scripts/step3_extract-posn_llama-3-1-70b-instruct.py --input "..." --output_dir "..." --first_page 7 --last_page 27
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
import requests
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||||
MODEL = "meta-llama/Meta-Llama-3.1-70B-Instruct"
|
||||
MODEL_TAG = "llama-3-1-70b" # used in output filenames
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAY = 10
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are a military historian extracting structured data from a WWII Canadian Army
|
||||
war diary file (Calgary Highlanders, September–October 1944).
|
||||
|
||||
You will receive the raw text of one page. Determine the page type and extract
|
||||
accordingly. Output valid JSON only. No prose before or after the JSON.
|
||||
|
||||
════════════════════════════════════════
|
||||
PAGE TYPE DETECTION
|
||||
════════════════════════════════════════
|
||||
|
||||
page_type must be one of:
|
||||
"diary" — HTML table with Place / Date / Hour / Summary / Remarks columns
|
||||
"message_form" — Army Form C2136 or similar signal/message form
|
||||
"movement_order" — Numbered movement or operation order
|
||||
"arty_table" — Artillery DF / SOS / fire task table
|
||||
"field_return" — Strength return (officers or other ranks)
|
||||
"isum" — Intelligence summary
|
||||
"patrol_report" — Patrol programme or patrol report
|
||||
"trace" — Map, sketch map, disposition diagram, or trace
|
||||
"admin" — Part I/II orders, nominal rolls, boilerplate instructions
|
||||
"other" — Anything that does not fit the above
|
||||
|
||||
════════════════════════════════════════
|
||||
OUTPUT SCHEMA — ALL PAGES
|
||||
════════════════════════════════════════
|
||||
|
||||
{
|
||||
"page_type": "diary | message_form | movement_order | ...",
|
||||
"date_warning": null or "explanation",
|
||||
|
||||
// ── DIARY PAGES ONLY ──────────────────────────────────────────────────────
|
||||
"narratives": [
|
||||
{
|
||||
"date": "1 Sep 44",
|
||||
"date_inferred": false,
|
||||
"date_span": null,
|
||||
"place": "Ste. Foy",
|
||||
"summary": "Full verbatim narrative text...",
|
||||
"weather": "Fair and warm",
|
||||
"no_change": false,
|
||||
"source_page": 7
|
||||
}
|
||||
],
|
||||
|
||||
// ── APPENDIX PAGES ONLY ───────────────────────────────────────────────────
|
||||
"appx_record": {
|
||||
"page_type": "message_form",
|
||||
"document_id": "GO-7",
|
||||
"date": "27 Sep 44",
|
||||
"time": "2205A",
|
||||
"from_unit": "HQ RCA 2 Cdn Inf Div",
|
||||
"to_units": ["5 CIB", "RHC", "CALG HIGHRS"],
|
||||
"subject": "DF Task Table No 35 — amendment",
|
||||
"summary": "One sentence description of what this document contains.",
|
||||
"units_mentioned": ["Calgary Highlanders", "RHC", "R de Mais", "5 Cdn Fd Regt"],
|
||||
"source_page": 12
|
||||
},
|
||||
|
||||
// ── ALL PAGES — positions found anywhere on the page ─────────────────────
|
||||
"positions": [
|
||||
{
|
||||
"date": "1 Sep 44",
|
||||
"date_inferred": false,
|
||||
"hour": null,
|
||||
"grid": "255330",
|
||||
"grid_inferred": false,
|
||||
"place_name": "Ste. Foy",
|
||||
"sheet_ref": "Sheet 861",
|
||||
"category": "END_OF_DAY",
|
||||
"subunit": null,
|
||||
"friendly_unit": null,
|
||||
"no_movement": false,
|
||||
"confidence": "high",
|
||||
"context": "arrived in the little village of Ste. Foy east of Longueville",
|
||||
"source_page": 7
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
For diary pages: populate "narratives" and "positions". Set "appx_record" to null.
|
||||
For appendix pages: populate "appx_record" and "positions". Set "narratives" to [].
|
||||
For blank pages: set all arrays empty, appx_record null.
|
||||
|
||||
════════════════════════════════════════
|
||||
DIARY PAGE RULES (page_type = "diary")
|
||||
════════════════════════════════════════
|
||||
|
||||
Each <tr> in the HTML table is one diary entry. Extract:
|
||||
|
||||
narratives:
|
||||
date — from Date <td>. Join multiline ("1 Sep\n44" → "1 Sep 44").
|
||||
date_span — if entry covers multiple dates ("6-8 Sep"), set to "6-8 Sep 44",
|
||||
date to the first date.
|
||||
date_inferred — true if date carried forward from a previous row.
|
||||
place — primary place name from Place <td>, excluding grid/sheet refs.
|
||||
summary — full verbatim text from Summary <td>. Nothing omitted.
|
||||
weather — weather note if present on the page. null if absent.
|
||||
no_change — true if entry states no change from previous day.
|
||||
|
||||
positions (cast a wide net — extract from BOTH Place <td> AND Summary <td>):
|
||||
PLACE COLUMN → category END_OF_DAY (last entry per date) or UNIT_MOVEMENT.
|
||||
SUMMARY TEXT → category SUBUNIT, FRIENDLY, ENEMY, PATROL, or MISC.
|
||||
Only ONE position per date may have category END_OF_DAY.
|
||||
|
||||
════════════════════════════════════════
|
||||
APPENDIX PAGE RULES (all other page_types)
|
||||
════════════════════════════════════════
|
||||
|
||||
appx_record:
|
||||
document_id — order/form number if visible (e.g. "Mov Order No 5",
|
||||
"ISUM No 45", "GO-3", "DF Task Table No 35"). null if absent.
|
||||
date — date of the document. null if not determinable.
|
||||
time — time or date-time group if present (e.g. "272205A"). null if absent.
|
||||
from_unit — originating unit/HQ. null if absent.
|
||||
to_units — list of addressees. [] if absent.
|
||||
subject — subject line or a one-sentence description of purpose.
|
||||
summary — one to three sentence plain-English summary of what this
|
||||
document records or orders. No verbatim transcription.
|
||||
units_mentioned — all unit names appearing anywhere on the page.
|
||||
|
||||
positions (extract ALL grid references and named locations):
|
||||
Extract every grid reference (4-digit, 6-digit, 8-digit) and every named
|
||||
location from the entire page, regardless of context.
|
||||
Use the same position schema as diary pages.
|
||||
category:
|
||||
UNIT_MOVEMENT — if associated with the Calgary Highlanders' own movement
|
||||
FRIENDLY — if associated with another Allied unit
|
||||
ENEMY — if associated with enemy forces
|
||||
PATROL — if from a patrol programme or patrol report
|
||||
DF_TASK — if from an DF/SOS artillery task table
|
||||
MISC — everything else (objectives, named features, route points)
|
||||
date — take from document date if not stated per-position. null if unknown.
|
||||
confidence:
|
||||
"high" — grid explicitly written next to a place name
|
||||
"medium" — grid present but context unclear
|
||||
"low" — place name only, no grid; or grid with no place name context
|
||||
|
||||
════════════════════════════════════════
|
||||
GRID REFERENCE RULES (all page types)
|
||||
════════════════════════════════════════
|
||||
|
||||
grid: digits only, no prefix, no punctuation.
|
||||
"MR 2468" → "2468" "GR 442891" → "442891" "MR 24.68" → "2468"
|
||||
If the diary gives a 4-digit grid, expand: "2553" → "255535". Set grid_inferred = true.
|
||||
Do not guess or expand 6-digit grids. Set grid_inferred = false.
|
||||
8-digit grids: keep first 3 + last 3 digits → "44289100" → "442891".
|
||||
5-digit grids: set grid = null, confidence = "low", note in context.
|
||||
If no grid is available, set grid to null and record place_name instead.
|
||||
|
||||
sheet_ref: carry forward the most recent sheet reference seen on the page.
|
||||
hour: HHMM format if stated ("0930"). null if absent.
|
||||
Ignore Hour <td> values that look like years ("44") or grid refs.
|
||||
|
||||
Output ONLY the JSON object. No preamble. No explanation. No trailing text.\
|
||||
"""
|
||||
|
||||
|
||||
def make_user_prompt(page_num: int, page_text: str) -> str:
|
||||
return (
|
||||
f"This is page {page_num} of the war diary file. "
|
||||
f"Determine the page type, then extract accordingly.\n\n"
|
||||
f"PAGE TEXT:\n{page_text}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MD page parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_md_pages(md_text: str) -> dict[int, str]:
|
||||
pages = {}
|
||||
page_re = re.compile(r'^## Page (\d+)\s*$', re.MULTILINE)
|
||||
matches = list(page_re.finditer(md_text))
|
||||
for i, m in enumerate(matches):
|
||||
page_num = int(m.group(1))
|
||||
start = m.end()
|
||||
end = matches[i + 1].start() if i + 1 < len(matches) else len(md_text)
|
||||
pages[page_num] = md_text[start:end].strip()
|
||||
return pages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON response parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_json_response(raw: str, page_num: int) -> dict | None:
|
||||
cleaned = raw.strip()
|
||||
cleaned = re.sub(r'^```json\s*', '', cleaned, flags=re.IGNORECASE)
|
||||
cleaned = re.sub(r'^```\s*', '', cleaned)
|
||||
cleaned = re.sub(r'\s*```$', '', cleaned)
|
||||
cleaned = cleaned.strip()
|
||||
try:
|
||||
return json.loads(cleaned)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f" WARNING: JSON parse failed for page {page_num}: {e}")
|
||||
print(f" Raw response (first 300 chars): {raw[:300]}")
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API call — DeepInfra (OpenAI-compatible)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def extract_page(page_num: int, page_text: str, api_key: str) -> dict | None:
|
||||
"""Send one page to Llama via DeepInfra and return parsed extraction result."""
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": make_user_prompt(page_num, page_text)},
|
||||
],
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(
|
||||
DEEPINFRA_API_URL,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# DeepInfra returns usage in the same response
|
||||
usage = data.get("usage", {})
|
||||
input_t = usage.get("prompt_tokens", 0)
|
||||
output_t = usage.get("completion_tokens", 0)
|
||||
print(f" tokens: {input_t} in / {output_t} out", end=" ")
|
||||
|
||||
raw_content = data["choices"][0]["message"]["content"]
|
||||
return parse_json_response(raw_content, page_num)
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if hasattr(e, 'response') and e.response is not None:
|
||||
print(f" Response: {e.response.text[:300]}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" Error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-of-day assignment (diary positions only)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def assign_end_of_day(positions: list[dict]) -> list[dict]:
|
||||
from collections import defaultdict
|
||||
date_groups: dict[str, list[int]] = defaultdict(list)
|
||||
for i, p in enumerate(positions):
|
||||
if p.get('date'):
|
||||
date_groups[p['date']].append(i)
|
||||
|
||||
for date, indices in date_groups.items():
|
||||
eod_indices = [i for i in indices if positions[i].get('category') == 'END_OF_DAY']
|
||||
if len(eod_indices) == 1:
|
||||
continue
|
||||
if len(eod_indices) > 1:
|
||||
for i in eod_indices[:-1]:
|
||||
positions[i]['category'] = 'UNIT_MOVEMENT'
|
||||
continue
|
||||
with_grid = [i for i in indices if positions[i].get('grid')]
|
||||
unit_movement = [i for i in with_grid if positions[i].get('category') == 'UNIT_MOVEMENT']
|
||||
if unit_movement:
|
||||
positions[unit_movement[-1]]['category'] = 'END_OF_DAY'
|
||||
elif with_grid:
|
||||
positions[with_grid[-1]]['category'] = 'END_OF_DAY'
|
||||
return positions
|
||||
|
||||
|
||||
def post_process_positions(positions: list[dict]) -> list[dict]:
|
||||
for p in positions:
|
||||
raw = re.sub(r'[^0-9]', '', p.get('grid') or '')
|
||||
if len(raw) == 4:
|
||||
p['grid'] = raw[0:2] + '5' + raw[2:4] + '5'
|
||||
p['grid_inferred'] = True
|
||||
elif len(raw) == 6:
|
||||
p['grid'] = raw
|
||||
p['grid_inferred'] = False
|
||||
elif len(raw) == 8:
|
||||
p['grid'] = raw[0:3] + raw[4:7]
|
||||
p['grid_inferred'] = False
|
||||
elif len(raw) == 5:
|
||||
p['grid'] = None
|
||||
p['grid_inferred'] = False
|
||||
p['confidence'] = 'low'
|
||||
p['context'] = (p.get('context') or '') + ' [5-figure grid — needs human review]'
|
||||
else:
|
||||
p['grid'] = None
|
||||
p['grid_inferred'] = False
|
||||
p['confidence'] = 'low'
|
||||
return positions
|
||||
|
||||
|
||||
def dedup_positions(positions: list[dict]) -> list[dict]:
|
||||
seen = set()
|
||||
result = []
|
||||
for p in positions:
|
||||
key = (p.get('date'), p.get('grid'), (p.get('context') or '')[:60])
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
result.append(p)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main extraction loop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def extract_diary(md_path: Path, output_dir: Path, api_key: str,
|
||||
first_page: int = None, last_page: int = None):
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Extracting: {md_path.name}")
|
||||
if first_page or last_page:
|
||||
print(f"Pages: {first_page or 'start'} → {last_page or 'end'}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
stem = md_path.stem
|
||||
for suffix in ('_olmocr', '_claude-4-6', 'step1_'):
|
||||
stem = stem.replace(suffix, '')
|
||||
stem = stem.strip('_')
|
||||
|
||||
# Diary outputs
|
||||
narratives_out = output_dir / f"{stem}_{MODEL_TAG}_narratives.json"
|
||||
positions_out = output_dir / f"{stem}_{MODEL_TAG}_positions.json"
|
||||
# Appendix outputs
|
||||
appx_records_out = output_dir / f"{stem}_{MODEL_TAG}_appx_records.json"
|
||||
appx_posns_out = output_dir / f"{stem}_{MODEL_TAG}_appx_positions.json"
|
||||
# Checkpoint file — deleted on clean completion
|
||||
checkpoint_path = output_dir / f"{stem}_{MODEL_TAG}_checkpoint.json"
|
||||
|
||||
md_text = md_path.read_text(encoding='utf-8')
|
||||
all_pages = parse_md_pages(md_text)
|
||||
|
||||
if not all_pages:
|
||||
print(" ERROR: No pages found in MD file. Check file format.")
|
||||
return
|
||||
|
||||
page_nums = sorted(all_pages.keys())
|
||||
if first_page:
|
||||
page_nums = [p for p in page_nums if p >= first_page]
|
||||
if last_page:
|
||||
page_nums = [p for p in page_nums if p <= last_page]
|
||||
|
||||
# ── Resume from checkpoint if one exists ─────────────────────────────────
|
||||
all_narratives = []
|
||||
all_positions = []
|
||||
all_appx_records = []
|
||||
all_appx_posns = []
|
||||
skipped_pages = []
|
||||
error_pages = []
|
||||
resume_from = None
|
||||
|
||||
if checkpoint_path.exists():
|
||||
try:
|
||||
ckpt = json.loads(checkpoint_path.read_text(encoding='utf-8'))
|
||||
resume_from = ckpt.get('last_completed_page')
|
||||
all_narratives = ckpt.get('narratives', [])
|
||||
all_positions = ckpt.get('positions', [])
|
||||
all_appx_records = ckpt.get('appx_records', [])
|
||||
all_appx_posns = ckpt.get('appx_positions',[])
|
||||
skipped_pages = ckpt.get('skipped_pages', [])
|
||||
error_pages = ckpt.get('error_pages', [])
|
||||
print(f" RESUMING from checkpoint — last completed page: {resume_from}")
|
||||
page_nums = [p for p in page_nums if p > resume_from]
|
||||
print(f" Remaining pages: {len(page_nums)}")
|
||||
except Exception as e:
|
||||
print(f" WARNING: Could not read checkpoint ({e}) — starting from scratch.")
|
||||
|
||||
print(f" Found {len(all_pages)} total pages, processing {len(page_nums)}.")
|
||||
print(f" Model: {MODEL}")
|
||||
print(f" Diary outputs : {narratives_out.name} / {positions_out.name}")
|
||||
print(f" Appendix outputs: {appx_records_out.name} / {appx_posns_out.name}")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def save_checkpoint(last_page_done: int):
|
||||
checkpoint_path.write_text(
|
||||
json.dumps({
|
||||
'last_completed_page': last_page_done,
|
||||
'narratives': all_narratives,
|
||||
'positions': all_positions,
|
||||
'appx_records': all_appx_records,
|
||||
'appx_positions':all_appx_posns,
|
||||
'skipped_pages': skipped_pages,
|
||||
'error_pages': error_pages,
|
||||
}, ensure_ascii=False),
|
||||
encoding='utf-8'
|
||||
)
|
||||
|
||||
for i, page_num in enumerate(page_nums, start=1):
|
||||
page_text = all_pages[page_num]
|
||||
|
||||
if not page_text.strip() or page_text.strip() == '[BLANK]':
|
||||
print(f" Page {page_num} ({i}/{len(page_nums)})... skipped (blank)")
|
||||
skipped_pages.append(page_num)
|
||||
save_checkpoint(page_num)
|
||||
continue
|
||||
|
||||
print(f" Page {page_num} ({i}/{len(page_nums)})...", end=" ", flush=True)
|
||||
|
||||
result = extract_page(page_num, page_text, api_key)
|
||||
|
||||
if result is None:
|
||||
print("ERROR — extraction failed")
|
||||
error_pages.append(page_num)
|
||||
save_checkpoint(page_num)
|
||||
continue
|
||||
|
||||
page_type = result.get('page_type', 'other')
|
||||
date_warning = result.get('date_warning')
|
||||
narratives = result.get('narratives', []) or []
|
||||
positions = result.get('positions', []) or []
|
||||
appx_record = result.get('appx_record')
|
||||
|
||||
for n in narratives:
|
||||
n['source_page'] = page_num
|
||||
for p in positions:
|
||||
p['source_page'] = page_num
|
||||
if appx_record:
|
||||
appx_record['source_page'] = page_num
|
||||
|
||||
if page_type == 'diary':
|
||||
all_narratives.extend(narratives)
|
||||
all_positions.extend(positions)
|
||||
parts = [f"diary — {len(narratives)} narrative(s), {len(positions)} position(s)"]
|
||||
else:
|
||||
if appx_record:
|
||||
all_appx_records.append(appx_record)
|
||||
all_appx_posns.extend(positions)
|
||||
doc_id = (appx_record or {}).get('document_id') or ''
|
||||
parts = [f"{page_type}{' — ' + doc_id if doc_id else ''} — {len(positions)} position(s)"]
|
||||
|
||||
if date_warning:
|
||||
parts.append(f"DATE WARNING: {date_warning}")
|
||||
print(", ".join(parts))
|
||||
|
||||
# Save checkpoint after every successfully processed page
|
||||
save_checkpoint(page_num)
|
||||
|
||||
if i < len(page_nums):
|
||||
time.sleep(0.5)
|
||||
|
||||
# Post-process diary positions
|
||||
print(f"\n Running end-of-day assignment (diary)...")
|
||||
all_positions = assign_end_of_day(all_positions)
|
||||
all_positions = post_process_positions(all_positions)
|
||||
all_positions = dedup_positions(all_positions)
|
||||
|
||||
# Post-process appendix positions
|
||||
all_appx_posns = post_process_positions(all_appx_posns)
|
||||
all_appx_posns = dedup_positions(all_appx_posns)
|
||||
|
||||
# Write final outputs
|
||||
narratives_out.write_text(
|
||||
json.dumps(all_narratives, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
positions_out.write_text(
|
||||
json.dumps(all_positions, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
appx_records_out.write_text(
|
||||
json.dumps(all_appx_records, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
appx_posns_out.write_text(
|
||||
json.dumps(all_appx_posns, indent=2, ensure_ascii=False), encoding='utf-8')
|
||||
|
||||
# Clean up checkpoint — run completed successfully
|
||||
if checkpoint_path.exists():
|
||||
checkpoint_path.unlink()
|
||||
print(f" Checkpoint removed.")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Complete.")
|
||||
print(f" Diary narratives : {len(all_narratives)} entries → {narratives_out.name}")
|
||||
print(f" Diary positions : {len(all_positions)} records → {positions_out.name}")
|
||||
print(f" Appendix records : {len(all_appx_records)} documents → {appx_records_out.name}")
|
||||
print(f" Appendix positions: {len(all_appx_posns)} records → {appx_posns_out.name}")
|
||||
if skipped_pages:
|
||||
print(f" Skipped : pages {skipped_pages} (blank)")
|
||||
if error_pages:
|
||||
print(f" ERRORS : pages {error_pages} — re-run with --first_page/--last_page to retry")
|
||||
|
||||
warnings = [(n.get('source_page'), n.get('date'), n.get('date_warning'))
|
||||
for n in all_narratives if n.get('date_warning')]
|
||||
if warnings:
|
||||
print(f"\n DATE WARNINGS ({len(warnings)}):")
|
||||
for page, date, warning in warnings:
|
||||
print(f" Page {page} | date={date} | {warning}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Extract narratives, positions, and appendix records from OCR .md files using Llama 3.1 70B via DeepInfra."
|
||||
)
|
||||
parser.add_argument("--input", required=True,
|
||||
help="Path to the step1_*_olmocr.md (or _claude-4-6.md) file.")
|
||||
parser.add_argument("--output_dir", required=True,
|
||||
help="Folder to write output JSON files.")
|
||||
parser.add_argument("--first_page", type=int, default=None,
|
||||
help="First page to process (1-based).")
|
||||
parser.add_argument("--last_page", type=int, default=None,
|
||||
help="Last page to process (1-based).")
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = os.getenv("DEEPINFRA_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: DEEPINFRA_API_KEY not set in environment or .env file")
|
||||
sys.exit(1)
|
||||
|
||||
md_path = Path(args.input)
|
||||
if not md_path.exists():
|
||||
print(f"Error: input file not found: {md_path}")
|
||||
sys.exit(1)
|
||||
|
||||
extract_diary(
|
||||
md_path, Path(args.output_dir), api_key,
|
||||
first_page=args.first_page,
|
||||
last_page=args.last_page,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
297
scripts/step5_embed_wardiary.py
Normal file
297
scripts/step5_embed_wardiary.py
Normal file
@@ -0,0 +1,297 @@
|
||||
"""
|
||||
step5_embed_wardiary.py
|
||||
───────────────────────
|
||||
Reads step3/step4 JSON output files from a TestRun (or any outputs folder) and
|
||||
embeds them into the pgvector document_chunks table.
|
||||
|
||||
Sources ingested:
|
||||
*_narratives.json → chunk_type=narrative (one chunk per diary entry)
|
||||
*_appx_records.json → chunk_type=appx_record (one chunk per document page)
|
||||
|
||||
Positions JSON (*_positions.json, *_appx_positions.json) are NOT embedded here —
|
||||
they belong in a PostGIS geometry table, not a vector search table.
|
||||
|
||||
Usage:
|
||||
python scripts/step5_embed_wardiary.py --input_dir TestRun/outputs --nationality canadian
|
||||
python scripts/step5_embed_wardiary.py --input_dir TestRun/outputs --nationality canadian --dry_run
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
from openai import OpenAI
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ── Embedding config (same as GoogleDocumentOCR/embed_and_store.py) ──────────
|
||||
EMBEDDING_MODEL = "BAAI/bge-base-en-v1.5"
|
||||
EMBEDDING_DIMS = 768
|
||||
BATCH_SIZE = 100
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("DEEPINFRA_API_KEY"),
|
||||
base_url="https://api.deepinfra.com/v1/openai",
|
||||
)
|
||||
|
||||
# ── DB (reuse GoogleDocumentOCR/db.py) ───────────────────────────────────────
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "GoogleDocumentOCR"))
|
||||
from db import init_db, get_conn
|
||||
from psycopg2.extras import execute_values
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Embedding
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||||
all_embeddings = []
|
||||
for i in range(0, len(texts), BATCH_SIZE):
|
||||
batch = texts[i:i + BATCH_SIZE]
|
||||
response = client.embeddings.create(input=batch, model=EMBEDDING_MODEL)
|
||||
all_embeddings.extend([item.embedding for item in response.data])
|
||||
return all_embeddings
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# DB insert (extended schema with chunk_type, date_entry, unit)
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def init_extended_schema():
|
||||
"""Add chunk_type, date_entry, unit columns if they don't exist yet."""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
for col, typedef in [
|
||||
("chunk_type", "TEXT DEFAULT 'unknown'"),
|
||||
("date_entry", "TEXT"),
|
||||
("unit", "TEXT"),
|
||||
]:
|
||||
cur.execute(f"ALTER TABLE document_chunks ADD COLUMN IF NOT EXISTS {col} {typedef};")
|
||||
conn.commit()
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def upsert_rows(rows: list[dict]):
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
execute_values(
|
||||
cur,
|
||||
"""
|
||||
INSERT INTO document_chunks
|
||||
(filename, page_num, chunk_index, text, embedding,
|
||||
nationality, corpus, chunk_type, date_entry, unit)
|
||||
VALUES %s
|
||||
""",
|
||||
[
|
||||
(
|
||||
r["filename"], r["page_num"], r["chunk_index"], r["text"], r["embedding"],
|
||||
r["nationality"], r["corpus"], r["chunk_type"],
|
||||
r.get("date_entry"), r.get("unit"),
|
||||
)
|
||||
for r in rows
|
||||
],
|
||||
template="(%s, %s, %s, %s, %s::vector, %s, %s, %s, %s, %s)"
|
||||
)
|
||||
conn.commit()
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Text builders — turn each JSON record into a single embeddable string
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def narrative_to_text(entry: dict) -> str:
|
||||
parts = []
|
||||
if entry.get("date"):
|
||||
parts.append(f"Date: {entry['date']}")
|
||||
if entry.get("place"):
|
||||
parts.append(f"Location: {entry['place']}")
|
||||
if entry.get("weather"):
|
||||
parts.append(f"Weather: {entry['weather']}")
|
||||
if entry.get("summary"):
|
||||
parts.append(entry["summary"])
|
||||
return "\n".join(parts).strip()
|
||||
|
||||
|
||||
def appx_record_to_text(record: dict) -> str:
|
||||
parts = []
|
||||
if record.get("page_type"):
|
||||
parts.append(f"Document type: {record['page_type']}")
|
||||
if record.get("document_id"):
|
||||
parts.append(f"Document ID: {record['document_id']}")
|
||||
if record.get("date"):
|
||||
parts.append(f"Date: {record['date']}")
|
||||
if record.get("time"):
|
||||
parts.append(f"Time: {record['time']}")
|
||||
if record.get("from_unit"):
|
||||
parts.append(f"From: {record['from_unit']}")
|
||||
if record.get("to_units"):
|
||||
parts.append(f"To: {', '.join(record['to_units'])}")
|
||||
if record.get("subject"):
|
||||
parts.append(f"Subject: {record['subject']}")
|
||||
if record.get("summary"):
|
||||
parts.append(record["summary"])
|
||||
if record.get("units_mentioned"):
|
||||
parts.append(f"Units mentioned: {', '.join(record['units_mentioned'])}")
|
||||
return "\n".join(parts).strip()
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Ingest functions
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def ingest_narratives(json_path: Path, nationality: str, unit: str, dry_run: bool):
|
||||
data = json.loads(json_path.read_text(encoding="utf-8"))
|
||||
print(f" {json_path.name}: {len(data)} narrative entries")
|
||||
|
||||
chunks = []
|
||||
for idx, entry in enumerate(data):
|
||||
text = narrative_to_text(entry)
|
||||
if not text:
|
||||
continue
|
||||
chunks.append({
|
||||
"filename": json_path.name,
|
||||
"page_num": entry.get("source_page"),
|
||||
"chunk_index": idx,
|
||||
"text": text,
|
||||
"nationality": nationality,
|
||||
"corpus": "war_diary_narrative",
|
||||
"chunk_type": "narrative",
|
||||
"date_entry": entry.get("date"),
|
||||
"unit": unit,
|
||||
})
|
||||
|
||||
if not chunks:
|
||||
print(" No embeddable entries.")
|
||||
return
|
||||
|
||||
if dry_run:
|
||||
print(f" DRY RUN — would embed+store {len(chunks)} chunks")
|
||||
print(f" Sample: {chunks[0]['text'][:120]}...")
|
||||
return
|
||||
|
||||
print(f" Embedding {len(chunks)} chunks...")
|
||||
embeddings = embed_texts([c["text"] for c in chunks])
|
||||
for c, emb in zip(chunks, embeddings):
|
||||
c["embedding"] = emb
|
||||
|
||||
upsert_rows(chunks)
|
||||
print(f" ✓ Stored {len(chunks)} narrative chunks")
|
||||
|
||||
|
||||
def ingest_appx_records(json_path: Path, nationality: str, unit: str, dry_run: bool):
|
||||
data = json.loads(json_path.read_text(encoding="utf-8"))
|
||||
print(f" {json_path.name}: {len(data)} appendix records")
|
||||
|
||||
chunks = []
|
||||
for idx, record in enumerate(data):
|
||||
text = appx_record_to_text(record)
|
||||
if not text:
|
||||
continue
|
||||
chunks.append({
|
||||
"filename": json_path.name,
|
||||
"page_num": record.get("source_page"),
|
||||
"chunk_index": idx,
|
||||
"text": text,
|
||||
"nationality": nationality,
|
||||
"corpus": "war_diary_appendix",
|
||||
"chunk_type": "appx_record",
|
||||
"date_entry": record.get("date"),
|
||||
"unit": unit,
|
||||
})
|
||||
|
||||
if not chunks:
|
||||
print(" No embeddable records.")
|
||||
return
|
||||
|
||||
if dry_run:
|
||||
print(f" DRY RUN — would embed+store {len(chunks)} chunks")
|
||||
print(f" Sample: {chunks[0]['text'][:120]}...")
|
||||
return
|
||||
|
||||
print(f" Embedding {len(chunks)} chunks...")
|
||||
embeddings = embed_texts([c["text"] for c in chunks])
|
||||
for c, emb in zip(chunks, embeddings):
|
||||
c["embedding"] = emb
|
||||
|
||||
upsert_rows(chunks)
|
||||
print(f" ✓ Stored {len(chunks)} appendix record chunks")
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Unit name extraction from filename
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
UNIT_PATTERNS = {
|
||||
"calgary-highlanders": "Calgary Highlanders",
|
||||
"calgary_highlanders": "Calgary Highlanders",
|
||||
"blackwatch": "Black Watch (RHC)",
|
||||
"black-watch": "Black Watch (RHC)",
|
||||
"5cib": "5 Canadian Infantry Brigade",
|
||||
}
|
||||
|
||||
def extract_unit(filename: str) -> str:
|
||||
lower = filename.lower()
|
||||
for key, name in UNIT_PATTERNS.items():
|
||||
if key in lower:
|
||||
return name
|
||||
return "unknown"
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Main
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Embed war diary JSON outputs into pgvector")
|
||||
parser.add_argument("--input_dir", required=True, help="Folder containing step3/step4 JSON files")
|
||||
parser.add_argument("--nationality", default="canadian", help="Nationality tag (default: canadian)")
|
||||
parser.add_argument("--dry_run", action="store_true", help="Print what would be done without writing to DB")
|
||||
args = parser.parse_args()
|
||||
|
||||
input_dir = Path(args.input_dir)
|
||||
if not input_dir.exists():
|
||||
print(f"ERROR: input_dir not found: {input_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
if not args.dry_run:
|
||||
init_db()
|
||||
init_extended_schema()
|
||||
|
||||
# Find all JSON files that are narratives or appx_records (skip positions)
|
||||
narrative_files = sorted(input_dir.glob("*_narratives.json"))
|
||||
appx_record_files = sorted(input_dir.glob("*_appx_records.json"))
|
||||
|
||||
skipped = [f.name for f in input_dir.glob("*.json")
|
||||
if f not in narrative_files + appx_record_files]
|
||||
|
||||
print(f"\n{'DRY RUN — ' if args.dry_run else ''}War Diary Embedding Ingest")
|
||||
print(f"Input dir : {input_dir}")
|
||||
print(f"Nationality: {args.nationality}")
|
||||
print(f"Narratives : {len(narrative_files)} file(s)")
|
||||
print(f"Appx records: {len(appx_record_files)} file(s)")
|
||||
if skipped:
|
||||
print(f"Skipping (positions/other): {', '.join(skipped)}\n")
|
||||
|
||||
for f in narrative_files:
|
||||
unit = extract_unit(f.name)
|
||||
print(f"\n[narrative] {f.name} -> unit: {unit}")
|
||||
ingest_narratives(f, args.nationality, unit, args.dry_run)
|
||||
|
||||
for f in appx_record_files:
|
||||
unit = extract_unit(f.name)
|
||||
print(f"\n[appx_record] {f.name} -> unit: {unit}")
|
||||
ingest_appx_records(f, args.nationality, unit, args.dry_run)
|
||||
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
|
||||
285
scripts/step6_generate_report.py
Normal file
285
scripts/step6_generate_report.py
Normal file
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
step6_generate_report.py
|
||||
─────────────────────────────────────────────────────────────────────────────
|
||||
Generates a Tier-3 descendant report by combining:
|
||||
1. Semantic search against the pgvector DB (war diary chunks)
|
||||
2. The prompt template from prompts/v2-tier3-report.md
|
||||
3. Claude (Anthropic) for the final report generation
|
||||
|
||||
Usage
|
||||
─────
|
||||
# Default — uses the Bloggins placeholder soldier
|
||||
python scripts/step6_generate_report.py
|
||||
|
||||
# Real soldier
|
||||
python scripts/step6_generate_report.py \
|
||||
--name "Pte. John Smith" \
|
||||
--unit "Calgary Highlanders" \
|
||||
--joined "mid September 1944" \
|
||||
--wounded "late October 1944" \
|
||||
--notes "Family knows he was in D Company. Was evacuated to England."
|
||||
|
||||
Output
|
||||
──────
|
||||
reports/<soldier-slug>_<timestamp>_report.md
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import textwrap
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# ── Paths ─────────────────────────────────────────────────────────────────────
|
||||
REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
PROMPT_FILE = REPO_ROOT / "prompts" / "v2-tier3-report.md"
|
||||
REPORTS_DIR = REPO_ROOT / "reports"
|
||||
|
||||
# ── Add GoogleDocumentOCR to path so we can import db + embed_and_store ───────
|
||||
sys.path.insert(0, str(REPO_ROOT / "GoogleDocumentOCR"))
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# ── Search queries to give broad coverage of the service window ────────────────
|
||||
SEARCH_QUERIES = [
|
||||
"Calgary Highlanders September 1944 operations movements France",
|
||||
"Calgary Highlanders October 1944 Scheldt fighting Holland",
|
||||
"Calgary Highlanders November 1944 casualties wounded Maas",
|
||||
"battalion advance attack objective company platoon",
|
||||
"reinforcements joined unit billets rest weather rations",
|
||||
"parades training administrative daily routine",
|
||||
"casualties killed wounded evacuated",
|
||||
"officers commanding lieutenant colonel major company commander",
|
||||
"civilian population towns villages liberated",
|
||||
"enemy German positions shelling mortars small arms fire",
|
||||
]
|
||||
|
||||
CHUNKS_PER_QUERY = 5 # results per query
|
||||
MAX_CONTEXT_CHARS = 80_000 # ~20k tokens — safe Claude window
|
||||
|
||||
|
||||
def retrieve_diary_chunks(nationality: str = "canadian",
|
||||
corpus: str = "war_diary_narrative") -> list[dict]:
|
||||
"""
|
||||
Run multiple semantic searches and return a deduplicated, sorted list
|
||||
of diary chunks.
|
||||
"""
|
||||
from embed_and_store import embed_texts
|
||||
from db import semantic_search
|
||||
|
||||
seen_ids = set()
|
||||
chunks = []
|
||||
|
||||
print(f" Running {len(SEARCH_QUERIES)} semantic searches...")
|
||||
for query in SEARCH_QUERIES:
|
||||
embedding = embed_texts([query])[0]
|
||||
rows = semantic_search(
|
||||
embedding,
|
||||
top_k=CHUNKS_PER_QUERY,
|
||||
nationality=nationality,
|
||||
corpus=corpus,
|
||||
)
|
||||
for row in rows:
|
||||
filename, page_num, chunk_index, text, nat, corp, similarity = row
|
||||
key = (filename, page_num, chunk_index)
|
||||
if key not in seen_ids:
|
||||
seen_ids.add(key)
|
||||
chunks.append({
|
||||
"filename": filename,
|
||||
"page_num": page_num,
|
||||
"chunk_index": chunk_index,
|
||||
"text": text,
|
||||
"similarity": similarity,
|
||||
})
|
||||
|
||||
# Sort by filename then page so the context reads chronologically
|
||||
chunks.sort(key=lambda c: (c["filename"], c["page_num"], c["chunk_index"]))
|
||||
print(f" Retrieved {len(chunks)} unique chunks after deduplication.")
|
||||
return chunks
|
||||
|
||||
|
||||
def build_context_block(chunks: list[dict], max_chars: int = MAX_CONTEXT_CHARS) -> str:
|
||||
"""
|
||||
Format chunks as a readable diary-extract block, trimmed to max_chars.
|
||||
"""
|
||||
lines = []
|
||||
current_file = None
|
||||
char_count = 0
|
||||
|
||||
for c in chunks:
|
||||
if c["filename"] != current_file:
|
||||
current_file = c["filename"]
|
||||
header = f"\n--- Source: {current_file} ---\n"
|
||||
lines.append(header)
|
||||
char_count += len(header)
|
||||
|
||||
entry = f"[page {c['page_num']}] {c['text']}\n\n"
|
||||
if char_count + len(entry) > max_chars:
|
||||
lines.append("\n[Context truncated — additional sources available in DB]\n")
|
||||
break
|
||||
lines.append(entry)
|
||||
char_count += len(entry)
|
||||
|
||||
return "".join(lines)
|
||||
|
||||
|
||||
def build_customer_situation(name: str, unit: str, joined: str,
|
||||
wounded: str, notes: str) -> str:
|
||||
knows = textwrap.dedent(f"""\
|
||||
The family knows:
|
||||
- Their relative, {name}, served with the {unit}.
|
||||
- They joined the battalion around {joined}.
|
||||
- They were wounded in action around {wounded}.
|
||||
- They were evacuated and did not return to the unit.
|
||||
""")
|
||||
|
||||
doesnt_know = textwrap.dedent("""\
|
||||
The family does NOT know:
|
||||
- The exact date they joined or were wounded.
|
||||
- Specific actions they were personally involved in.
|
||||
- Their company, platoon, or section.
|
||||
""")
|
||||
|
||||
if notes:
|
||||
knows += f"\nAdditional context provided by the family:\n{notes}\n"
|
||||
|
||||
return f"## The customer's situation\n{knows}\n{doesnt_know}"
|
||||
|
||||
|
||||
def load_prompt_template() -> str:
|
||||
return PROMPT_FILE.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def assemble_full_prompt(template: str, customer_block: str,
|
||||
context_block: str) -> str:
|
||||
"""
|
||||
Replace the customer situation section in the template with the real
|
||||
soldier info, then append the diary context before the output instructions.
|
||||
"""
|
||||
# Strip the hardcoded customer section and inject the real one
|
||||
start_marker = "## The customer's situation"
|
||||
end_marker = "## What to produce"
|
||||
|
||||
if start_marker in template and end_marker in template:
|
||||
before = template[:template.index(start_marker)]
|
||||
after = template[template.index(end_marker):]
|
||||
template = before + customer_block + "\n\n" + after
|
||||
|
||||
diary_section = (
|
||||
"## War diary source material\n"
|
||||
"The following excerpts are drawn from the pgvector database of "
|
||||
"OCR-processed war diaries. Use them as your primary source.\n\n"
|
||||
+ context_block
|
||||
+ "\n\n"
|
||||
)
|
||||
|
||||
# Insert diary context just before "## What to produce"
|
||||
if end_marker in template:
|
||||
idx = template.index(end_marker)
|
||||
template = template[:idx] + diary_section + template[idx:]
|
||||
else:
|
||||
template += "\n\n" + diary_section
|
||||
|
||||
return template
|
||||
|
||||
|
||||
def generate_report_with_claude(prompt: str, soldier_name: str) -> str:
|
||||
import anthropic
|
||||
|
||||
client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
|
||||
|
||||
print(" Sending to Claude... (this may take 30–60 seconds)")
|
||||
message = client.messages.create(
|
||||
model="claude-sonnet-4-5",
|
||||
max_tokens=4096,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
)
|
||||
return message.content[0].text
|
||||
|
||||
|
||||
def save_report(text: str, soldier_name: str) -> Path:
|
||||
REPORTS_DIR.mkdir(exist_ok=True)
|
||||
slug = soldier_name.lower().replace(" ", "_").replace(".", "").replace(",", "")
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
out_path = REPORTS_DIR / f"{slug}_{timestamp}_report.md"
|
||||
out_path.write_text(text, encoding="utf-8")
|
||||
return out_path
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a Tier-3 descendant report from war diary pgvector DB"
|
||||
)
|
||||
parser.add_argument("--name", default="Pte. Bill Bloggins",
|
||||
help="Soldier's name (default: placeholder Bloggins)")
|
||||
parser.add_argument("--unit", default="Calgary Highlanders",
|
||||
help="Unit name")
|
||||
parser.add_argument("--joined", default="mid September 1944",
|
||||
help="Approximate join date (plain English)")
|
||||
parser.add_argument("--wounded", default="late October 1944",
|
||||
help="Approximate wound/end date (plain English)")
|
||||
parser.add_argument("--notes", default="",
|
||||
help="Any extra family context (optional)")
|
||||
parser.add_argument("--nationality", default="canadian",
|
||||
help="DB nationality filter (default: canadian)")
|
||||
parser.add_argument("--corpus", default="war_diary_narrative",
|
||||
help="DB corpus filter (default: war_diary_narrative)")
|
||||
parser.add_argument("--top-k", type=int, default=5,
|
||||
help="Results per search query (default: 5)")
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"\n=== Generating report for: {args.name} ===")
|
||||
print(f" Unit : {args.unit}")
|
||||
print(f" Joined : {args.joined}")
|
||||
print(f" Wounded : {args.wounded}")
|
||||
if args.notes:
|
||||
print(f" Notes : {args.notes}")
|
||||
print()
|
||||
|
||||
# 1. Retrieve diary chunks from pgvector
|
||||
print("[1/4] Retrieving diary chunks from pgvector DB...")
|
||||
chunks = retrieve_diary_chunks(
|
||||
nationality=args.nationality,
|
||||
corpus=args.corpus,
|
||||
)
|
||||
if not chunks:
|
||||
print("ERROR: No chunks found. Check DB connection and filters.")
|
||||
sys.exit(1)
|
||||
|
||||
# 2. Build context block
|
||||
print("[2/4] Assembling context...")
|
||||
context_block = build_context_block(chunks)
|
||||
print(f" Context: {len(context_block):,} characters from {len(chunks)} chunks")
|
||||
|
||||
# 3. Assemble final prompt
|
||||
print("[3/4] Building prompt from v2-tier3-report.md...")
|
||||
template = load_prompt_template()
|
||||
customer_block = build_customer_situation(
|
||||
name=args.name, unit=args.unit,
|
||||
joined=args.joined, wounded=args.wounded,
|
||||
notes=args.notes,
|
||||
)
|
||||
full_prompt = assemble_full_prompt(template, customer_block, context_block)
|
||||
print(f" Total prompt: {len(full_prompt):,} characters")
|
||||
|
||||
# 4. Generate with Claude
|
||||
print("[4/4] Generating report with Claude...")
|
||||
report_text = generate_report_with_claude(full_prompt, args.name)
|
||||
|
||||
# 5. Save
|
||||
out_path = save_report(report_text, args.name)
|
||||
print(f"\n✅ Report saved to: {out_path}")
|
||||
print("-" * 60)
|
||||
print(report_text[:500] + "...")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user