import json from pathlib import Path import pytest from curatorbot import agent, backtest, gazetteer, transcript as tx ROOT = Path(__file__).resolve().parent.parent DEMO = json.loads((ROOT / "data/interviews/DEMO-0001/brief.json").read_text(encoding="utf-8")) SEGS = DEMO["testimony"]["segments"] def test_parse_qa_text_handles_times_roles_and_continuations(): segs = tx.parse_qa_text("[00:01:02] Q: Where were you born?\n[1:05] A: In Lodz.\nWe moved later.\nINT: And then?") assert [s["role"] for s in segs] == ["Q", "A", "Q"] assert segs[0]["start"] == "00:01:02" and segs[1]["start"] == "00:01:05" assert segs[1]["text"] == "In Lodz. We moved later." def test_from_asr_uses_question_rate_to_find_interviewer(): chunks = [ {"start": 0, "end": 3, "text": "Where were you born?", "speaker": "S1"}, {"start": 3.5, "end": 9, "text": "In Lodz, in 1928.", "speaker": "S2"}, {"start": 9.5, "end": 12, "text": "We lived near the market.", "speaker": "S2"}, {"start": 13, "end": 15, "text": "And your father?", "speaker": "S1"}, ] segs = tx.from_asr(chunks) assert [s["role"] for s in segs] == ["Q", "A", "Q"] assert segs[1]["text"] == "In Lodz, in 1928. We lived near the market." def test_gazetteer_reviews_period_names(): ms =gazetteer.review_mentions("My uncle lived in Lemberg and Lwow.", [{"text": "Lemberg", "label": "PERSON", "start_char": 21, "end_char": 28}], "s1") by = {m["surface"]: m["action"] for m in ms} assert by == {"Lemberg": "relabelled", "Lwow": "added"} assert gazetteer.resolve_name("Łódź ghetto (Bałuty)") in {"lodz", "baluty"} assert gazetteer.resolve_name("American Jewish committee") is None # 'America' must not match inside 'American' def test_grounding_accepts_verbatim_and_repairs_wrong_segment(): ok, where = agent.grounded(SEGS, "s02", "In May 1947 we sailed from Bremerhaven to New York.") assert ok and where == "s23" assert agent.grounded(SEGS, "s02", "She was liberated at Theresienstadt.") == (False, None) assert agent.grounded(SEGS, "s10", "we will survive this, children")[0] # matches the translation def test_validate_rejects_ungrounded_and_out_of_vocabulary(): raw = {"headings": [ {"heading": "Auschwitz (Concentration camp)", "category": "Corporate", "segId": "s13", "quote": "the train arrived at Auschwitz"}, {"heading": "Invented heading", "category": "Topical", "segId": "s13", "quote": "the train arrived at Auschwitz"}, {"heading": "Dachau (Concentration camp)", "category": "Corporate", "segId": "s13", "quote": "we were taken to Dachau"}, ], "claims": [{"field": "Birth", "donor": "1929", "verdict": "not_mentioned"}]} out = agent.validate(raw, SEGS, ("Auschwitz (Concentration camp)", "Dachau (Concentration camp)")) flags = [(h["grounded"], h["in_vocabulary"]) for h in out["headings"]] assert flags == [(True, True), (True, False), (False, True)] assert out["claims"][0]["grounded"] is True assert out["validation"] == {"verified": 2, "rejected": 2, "quotes_rejected": 1, "out_of_vocabulary": 1} def test_demo_brief_quotes_are_all_grounded(): for h in DEMO["agent"]["headings"]: if h.get("source") == "rule" and not h.get("quote"): continue # convention headings proposed by rule carry their reason, not a quote assert agent.grounded(SEGS, h["segId"], h["quote"])[0], h for c in DEMO["agent"]["claims"]: if c["quote"]: assert agent.grounded(SEGS, c["segId"], c["quote"])[0], c def test_backtest_scores_exact_and_head_matches(): r = backtest.score(["Auschwitz (Concentration camp)", "Holocaust survivors", "Partisans"], ["Auschwitz (Concentration camp)--Personal narratives", "Holocaust survivors.", "Ghettos, Jewish--Poland--Łódź"]) assert r["recall_exact"] == pytest.approx(1 / 3, abs=1e-3) assert r["recall_head"] == pytest.approx(2 / 3, abs=1e-3) assert r["missed"] == ["Ghettos, Jewish--Poland--Łódź"] def test_extract_json_tolerates_fences(): from curatorbot.llm import extract_json assert extract_json('Here:\n```json\n{"a": "b}"}\n```') == {"a": "b}"} def test_heading_matching_ignores_final_full_stop_and_case(): from curatorbot.agent import validate segs = [{"segId": "p001", "role": "A", "start": "", "text": "I escaped from Poland and made my way to Germany."}] vocab = ("Poland--Emigration and immigration.", "Antisemitism") raw = {"headings": [ {"heading": "Poland--Emigration and immigration", "segId": "p001", "quote": "I escaped from Poland", "confidence": 0.5}, {"heading": "antisemitism", "segId": "p001", "quote": "made my way to Germany", "confidence": 0.5}, {"heading": "Not a heading", "segId": "p001", "quote": "I escaped", "confidence": 0.5}, {"heading": "Antisemitism", "segId": "p001", "quote": "words that are not there", "confidence": 0.5}]} out = validate(raw, segs, vocab) h = out["headings"] assert h[0]["heading"] == "Poland--Emigration and immigration." and h[0]["in_vocabulary"] assert h[0]["heading_as_proposed"] == "Poland--Emigration and immigration" assert h[1]["heading"] == "Antisemitism" and h[1]["in_vocabulary"] assert h[2]["in_vocabulary"] is False and h[3]["grounded"] is False v = out["validation"] assert (v["verified"], v["quotes_rejected"], v["out_of_vocabulary"], v["rejected"]) == (3, 1, 1, 2) def test_revalidate_cli_rewrites_brief_and_backtest(tmp_path): import shutil from curatorbot.cli import main d = tmp_path / "DEMO-0001" shutil.copytree(Path(__file__).resolve().parent.parent / "data/interviews/DEMO-0001", d) cat = tmp_path / "catalogue.json" cat.write_text(json.dumps({"id": "DEMO-0001", "subjects": ["Holocaust survivors"]}), encoding="utf-8") assert main(["revalidate", str(d / "brief.json"), "--catalogue", str(cat)]) == 0 brief = json.loads((d / "brief.json").read_text(encoding="utf-8")) assert "out_of_vocabulary" in brief["agent"]["validation"] and brief["run"]["revalidated"]["before"] is not None assert json.loads((d / "backtest.json").read_text(encoding="utf-8"))["catalogued"] == 1 def test_form_headings_by_rule_follow_grounded_gender_only(): from curatorbot.agent import validate segs = [{"segId": "p001", "role": "A", "start": "", "text": "I was the only girl in my class."}] vocab = ("Women--Personal narratives.", "Men--Personal narratives.", "Holocaust survivors--Interviews.") out = validate({"narrator": {"gender": "female", "segId": "p001", "quote": "I was the only girl"}, "headings": []}, segs, vocab) rule = {h["heading"]: h for h in out["headings"] if h["source"] == "rule"} assert set(rule) == {"Women--Personal narratives.", "Holocaust survivors--Interviews."} assert rule["Women--Personal narratives."]["quote"] == "I was the only girl" # an unsupported gender claim gives no gendered heading out = validate({"narrator": {"gender": "male", "segId": "p001", "quote": "I was a boy"}, "headings": []}, segs, vocab) assert [h["heading"] for h in out["headings"]] == ["Holocaust survivors--Interviews."] # revalidating keeps one copy of each rule heading again = validate(out, segs, vocab) assert [h["heading"] for h in again["headings"]] == ["Holocaust survivors--Interviews."] def test_places_gazetteer_then_geonames_then_near(): from curatorbot import places assert places.place("Litzmannstadt ghetto")["key"] == "lodz" r = places.place("the Częstochowa ghetto") assert r["source"] == "geonames" and r["name"] == "Częstochowa" and r["precision"] == "place" n = places.place("Pomortsy", near="Lodz") assert n["precision"] == "near" and n["anchor"] == "lodz" and n["approximate"] assert places.place("Pomortsy") is None assert places.place("Palestine")["precision"] == "country" j = [{"from": "Radomsko", "to": "Pomortsy"}] found = {} places.resolve_journeys(j, found) assert j[0]["fromPrecision"] == "place" and j[0]["toPrecision"] == "unplaced" and j[0]["toKey"] is None assert places.coverage(j, [])["journeys_not_drawn"] == 1 def test_backtest_keeps_model_metrics_and_adds_rule_and_r_precision(): brief = {"agent": {"headings": [ {"heading": "A", "confidence": 0.9}, {"heading": "B", "confidence": 0.8}, {"heading": "X", "confidence": 0.7}, {"heading": "Women--Personal narratives.", "confidence": 1.0, "source": "rule"}]}} rec = {"id": "t", "subjects": ["A", "B", "C", "Women--Personal narratives."]} bt = backtest.score_brief(brief, rec, vocab=("A", "B", "C", "X", "Women--Personal narratives.")) assert bt["suggested"] == 3 and bt["recall_exact"] == 0.5 # model only, as before w = bt["with_rules"] assert w["r"] == 4 and w["r_precision"] == 0.75 and w["recall_exact"] == 0.75 and w["precision_exact"] == 0.75 assert bt["rule_headings"] == ["Women--Personal narratives."] def test_unnamed_and_phonetic_places(): from curatorbot import places assert places.is_unnamed("a farm in a village (not named)") and places.is_unnamed("the camp for Poles") assert not places.is_unnamed("Varka (ph)") and not places.is_unnamed("the Częstochowa ghetto") w = places.place("Varka (ph)") assert w["name"] == "Warka" and w["uncertain"] and "by ear" in w["note"] j = [{"from": "his home town (name not given in transcript)", "to": "Warsaw"}, {"from": "Warsaw", "to": "Garbutka (ph)"}, {"from": "Warsaw", "to": "Zzyzxq"}] found = {} places.resolve_journeys(j, found) assert [j[0]["fromPrecision"], j[1]["toPrecision"], j[2]["toPrecision"]] == ["unnamed", "uncertain", "unplaced"] cov = places.coverage(j, []) assert (cov["journeys_drawn"], cov["journeys_not_drawn_unnamed"], cov["journeys_not_drawn_not_found"]) == (1, 1, 1) def test_prompt_states_the_heading_limit(monkeypatch): monkeypatch.setenv("CURATORBOT_MAX_HEADINGS", "17") prompt = agent.build_prompt(SEGS[:2], {}, ("Antisemitism",)) assert "at most 17;" in prompt def test_compare_reports_wer_turns_and_place_names(): pytest.importorskip("jiwer") from curatorbot.compare import compare ref = [{"segId": "p001", "role": "Q", "text": "Where were you?"}, {"segId": "p002", "role": "A", "text": "We were taken from Radomsko to Częstochowa (ph), to the ghetto."}] hyp = [{"segId": "s1", "role": "Q", "text": "Where were you"}, {"segId": "s2", "role": "A", "text": "we were taken from Radomsko to Chestohova to the ghetto"}] r = compare(ref, hyp) assert 0 < r["wer"] < 0.3 and r["turns"]["automatic"] == {"Q": 1, "A": 1} assert r["places"]["reference"] == 2 and r["places"]["found"] == 1 and r["places"]["missed"] == ["Częstochowa"] def test_compare_place_names_stop_at_sentence_ends(): pytest.importorskip("jiwer") from curatorbot.compare import compare ref = [{"segId": "p001", "role": "A", "text": "We went to the Ghetto Warsaw. It was closed. Then Warsaw. There we stayed."}] hyp = [{"segId": "s1", "role": "A", "text": "we went to the Warsaw ghetto it was closed then Warsaw there we stayed"}] r = compare(ref, hyp) assert r["places"]["reference"] == 1 and r["places"]["found"] == 1 and r["places"]["missed"] == [] def test_backtest_ranks_rule_headings_after_the_model(): brief = {"agent": {"headings": [ {"heading": "A", "confidence": 0.6}, {"heading": "B", "confidence": 0.5}, {"heading": "Women--Personal narratives.", "confidence": 1.0, "source": "rule"}]}} bt = backtest.score_brief(brief, {"id": "t", "subjects": ["A", "B"]}, vocab=("A", "B", "Women--Personal narratives.")) assert bt["with_rules"]["r_precision"] == 1.0 # top 2 are the model's A and B, not the rule heading def test_inventory_reads_a_box_style_delivery(tmp_path): from curatorbot.intake import find_media, inventory, interview_id two = tmp_path / "Interview 2" / "nested" two.mkdir(parents=True) for n in (1, 2, 4): (two / f"RG-50.030.0001.0{n}.04.mp4").write_bytes(b"x") five = tmp_path / "Interview 5" five.mkdir() (five / "tape A.mov").write_bytes(b"x") (five / "donor form.pdf").write_bytes(b"x") (tmp_path / "Interview 10").mkdir() entries = inventory(tmp_path) assert [e["folder"] for e in entries] == ["Interview 2", "Interview 5", "Interview 10"] # natural order two_e, five_e, ten_e = entries assert two_e["id"] == "RG-50.030.0001-BOX" and len(two_e["media"]) == 3 assert any("Missing part(s): 3 of 4" in i for i in two_e["issues"]) assert five_e["id"] == "BOX-Interview-5" and five_e["other_files"] == ["donor form.pdf"] assert ten_e["media"] == [] and "empty" in ten_e["issues"][0] assert [f.name for f in find_media(tmp_path / "Interview 2")][0] == "RG-50.030.0001.01.04.mp4" assert interview_id([], "Interview 9") == "BOX-Interview-9" def test_transcribe_reports_progress(monkeypatch, tmp_path, capsys): import types from curatorbot import asr from curatorbot.progress import Progress class Seg: def __init__(self, t): self.start, self.end, self.text, self.avg_logprob, self.no_speech_prob = t, t + 9, f"line {t}", -0.2, 0.0 class FakeModel: def transcribe(self, path, **kw): info = types.SimpleNamespace(duration=100.0, language="en", language_probability=0.98) return (Seg(t) for t in range(0, 100, 10)), info monkeypatch.setattr(asr, "_model", lambda size, device: FakeModel()) monkeypatch.setattr(asr, "load_audio", lambda path: [0.0]) prog = Progress(tmp_path / "progress.json", every=0) res = asr.transcribe("x.wav", size="large-v3", device="cuda", per_segment_lid=False, progress=prog) out = capsys.readouterr().out assert len(res["chunks"]) == 10 assert "Transcribing" in out and "of audio" in out and "(100%)" in out state = json.loads((tmp_path / "progress.json").read_text()) assert state["stage"] == "Transcribing" and state["percent"] == 100.0 def test_load_audio_reads_wav_without_pyav(tmp_path): import wave import numpy as np from curatorbot import asr samples = (np.sin(np.arange(16000) / 10) * 10000).astype(np.int16) with wave.open(str(tmp_path / "a.wav"), "wb") as w: w.setnchannels(1); w.setsampwidth(2); w.setframerate(16000); w.writeframes(samples.tobytes()) audio = asr.load_audio(str(tmp_path / "a.wav")) assert audio.dtype == np.float32 and len(audio) == 16000 and abs(audio.max() - samples.max() / 32768) < 1e-6 def test_speech_model_failure_explains_gpu_problem(monkeypatch): from curatorbot import asr def boom(size, device): raise OSError("Could not load library libcudnn_ops.so.9") monkeypatch.setattr(asr, "_model", boom) with pytest.raises(RuntimeError, match="nvidia-cublas-cu12"): asr.transcribe("x.wav", device="cuda") def test_missing_cublas_during_transcription_gets_the_hint(monkeypatch): from curatorbot import asr class FakeModel: def transcribe(self, audio, **kw): raise RuntimeError("Library libcublas.so.12 is not found or cannot be loaded") monkeypatch.setattr(asr, "_model", lambda size, device: FakeModel()) monkeypatch.setattr(asr, "load_audio", lambda path: [0.0]) with pytest.raises(RuntimeError, match="libcublas.so.12 is not found.*pip install nvidia-cublas-cu12"): asr.transcribe("x.wav", device="cuda") def test_progress_waiting_reports_while_a_step_runs(tmp_path, capsys): import time from curatorbot.progress import Progress prog = Progress(tmp_path / "p.json", every=0.05) with prog.waiting("Drafting the brief", "model x"): time.sleep(0.18) out = capsys.readouterr().out assert "Drafting the brief: model x" in out and "still working" in out and "done in" in out def test_build_brief_reports_each_stage(monkeypatch, tmp_path, capsys): from curatorbot import annotate, pipeline from curatorbot.progress import Progress class FakeAnn: def annotate(self, text, include_events=True): return {"entities": [], "event_data": []} class FakeJx: def extract(self, text, file_id=None): return {"journeys": []} monkeypatch.setattr(annotate, "_annotator", lambda m: FakeAnn()) monkeypatch.setattr(annotate, "_journey_rules", lambda m: FakeJx()) segs = [{"segId": f"s{i}", "start": f"00:00:{i:02d}", "role": "A" if i % 2 else "Q", "text": "We lived in Lodz."} for i in range(6)] pipeline.build_brief(segs, interview_id="T", title="T", use_llm=False, progress=Progress(tmp_path / "p.json", every=0)) out = capsys.readouterr().out for stage in ("Loading the language model", "Finding places, people and events", "3 of 3 narrator turns (100%)", "Placing", "Brief ready"): assert stage in out def test_harvest_ehri_headings_builds_a_held_out_vocabulary(tmp_path, monkeypatch, capsys): import importlib.util spec = importlib.util.spec_from_file_location("harvest", Path(__file__).resolve().parent.parent / "scripts" / "harvest_ehri_headings.py") harvest = importlib.util.module_from_spec(spec); spec.loader.exec_module(harvest) ead = tmp_path / "c.xml" ead.write_text(""" RG-50.030 Collection level. irn1RG-50.030*0034Interview A Forced labor.Only in A.Germany. Gestapo Someone, A. RG-50.030*0099Interview B Forced laborŁódź (Poland) RG-50.030*0100Interview C Forced labor.Oral histories. """, encoding="utf-8") cat = tmp_path / "RG-50.030.0034" / "catalogue.json" cat.parent.mkdir() cat.write_text(json.dumps({"id": "RG-50.030.0034", "subjects": ["Forced labor.", "Only in A.", "Germany. Gestapo."]})) assert harvest.main([str(ead), "--exclude", str(cat), "--out", str(tmp_path / "out")]) == 0 vocab = [l for l in (tmp_path / "out" / "vocabulary.txt").read_text(encoding="utf-8").splitlines() if not l.startswith("#")] assert vocab == ["Forced labor.", "Oral histories.", "Łódź (Poland)"] # A's own headings gone; shared one kept out = capsys.readouterr().out assert "RG-50.030.0034: 1 of 3 catalogue headings available" in out and "not available: Only in A." in out recs = [json.loads(l) for l in (tmp_path / "out" / "records.jsonl").read_text(encoding="utf-8").splitlines()] assert recs[0]["id"] == "RG-50.030*0034" and recs[0]["people"] == ["Someone, A."] and len(recs) == 3 # matched by the record number in the catalogue URL, or by title, when the accession number is not in the export cat.write_text(json.dumps({"id": "RG-50.030.9999", "url": "https://collections.ushmm.org/search/catalog/irn1", "subjects": ["Only in A."]})) assert harvest.main([str(ead), "--exclude", str(cat), "--out", str(tmp_path / "o3")]) == 0 assert "excluded Interview A" in capsys.readouterr().out cat.write_text(json.dumps({"id": "X", "title": "Interview b", "subjects": []})) assert harvest.main([str(ead), "--exclude", str(cat), "--out", str(tmp_path / "o4")]) == 0 assert "excluded Interview B" in capsys.readouterr().out assert harvest.main([str(ead), "--out", str(tmp_path / "o2"), "--min-records", "2"]) == 0 assert "Forced labor." in (tmp_path / "o2" / "vocabulary.txt").read_text(encoding="utf-8").splitlines()[2:] def test_vocabulary_base_can_be_replaced(tmp_path, monkeypatch): from curatorbot import vocabulary f = tmp_path / "v.txt" f.write_text("# x\nOne.\nTwo.\n", encoding="utf-8") monkeypatch.setenv("CURATORBOT_VOCABULARY", str(f)) monkeypatch.delenv("CURATORBOT_VOCAB_EXTRA", raising=False) assert vocabulary.load() == ("One.", "Two.")