Spaces:
Running
Running
| # Implementation Plan: `tei-annotator` | |
| Prompt: | |
| > Design a Python library called `tei-annotator` for annotating text with TEI XML tags using a two-stage LLM pipeline. The library should: | |
| > | |
| > **Inputs:** | |
| > | |
| > - A text string to annotate (may already contain partial XML) | |
| > - An injected `call_fn: (str) -> str` for calling an arbitrary inference endpoint | |
| > - An `EndpointCapability` enum indicating whether the endpoint is plain text generation, JSON-constrained, or a native extraction model like GLiNER2 | |
| > - A `TEISchema` data structure describing a subset of TEI elements with descriptions, allowed attributes, and legal child elements | |
| > The source text must never be modified by any model. Provide a package structure, all key data structures, and a step-by-step execution flow. | |
| ## Package Structure | |
| ``` | |
| tei_annotator/ | |
| ├── __init__.py | |
| ├── models/ | |
| │ ├── schema.py # TEI element/attribute data structures | |
| │ └── spans.py # Span manifest data structures | |
| ├── detection/ | |
| │ └── gliner_detector.py # Optional local GLiNER first-pass span detection | |
| ├── chunking/ | |
| │ └── chunker.py # Overlap-aware text chunker, XML-safe boundaries | |
| ├── prompting/ | |
| │ ├── builder.py # Prompt assembly | |
| │ └── templates/ | |
| │ ├── text_gen.jinja2 # For plain text-generation endpoints | |
| │ └── json_enforced.jinja2 # For JSON-mode / constrained endpoints | |
| ├── inference/ | |
| │ └── endpoint.py # Endpoint wrapper + capability enum | |
| ├── postprocessing/ | |
| │ ├── resolver.py # Context-anchor → char offset resolution | |
| │ ├── validator.py # Span verification + schema validation | |
| │ └── injector.py # Deterministic XML construction | |
| └── pipeline.py # Top-level orchestration | |
| ``` | |
| --- | |
| ## Data Structures (`models/`) | |
| ```python | |
| # schema.py | |
| @dataclass | |
| class TEIAttribute: | |
| name: str # e.g. "ref", "type", "cert" | |
| description: str | |
| required: bool = False | |
| allowed_values: list[str] | None = None # None = free string | |
| @dataclass | |
| class TEIElement: | |
| tag: str # e.g. "persName" | |
| description: str # from TEI Guidelines | |
| allowed_children: list[str] # tags of legal child elements | |
| attributes: list[TEIAttribute] | |
| @dataclass | |
| class TEISchema: | |
| elements: list[TEIElement] | |
| # convenience lookup | |
| def get(self, tag: str) -> TEIElement | None: ... | |
| # spans.py | |
| @dataclass | |
| class SpanDescriptor: | |
| element: str | |
| text: str | |
| context: str # must contain text as substring | |
| attrs: dict[str, str] | |
| # always flat — nesting is inferred from offset containment in resolver/injector, | |
| # not emitted by the model (models produce unreliable nested trees) | |
| confidence: float | None = None # passed through from GLiNER | |
| @dataclass | |
| class ResolvedSpan: | |
| element: str | |
| start: int | |
| end: int | |
| attrs: dict[str, str] | |
| children: list["ResolvedSpan"] | |
| fuzzy_match: bool = False # flagged for human review | |
| ``` | |
| --- | |
| ## GLiNER Dependency (`detection/gliner_detector.py`) | |
| `gliner` is a regular package dependency declared in `pyproject.toml` and installed via `uv add gliner`. No manual setup step is needed. | |
| Model weights are fetched from HuggingFace Hub automatically on first use of `GLiNER.from_pretrained(model_id)` and cached in `~/.cache/huggingface/`. If the import fails at runtime (e.g. the optional extra was not installed), the module raises a standard `ImportError` with a clear message — no wrapper needed. | |
| Recommended models (specified as `gliner_model` parameter): | |
| - `urchade/gliner_medium-v2.1` — balanced, Apache 2.0 | |
| - `numind/NuNER_Zero` — stronger multi-word entities, MIT (default) | |
| - `knowledgator/gliner-multitask-large-v0.5` — adds relation extraction | |
| All models run on CPU; no GPU required. | |
| --- | |
| ## Endpoint Abstraction (`inference/endpoint.py`) | |
| ```python | |
| class EndpointCapability(Enum): | |
| TEXT_GENERATION = "text_generation" # plain LLM, JSON via prompt only | |
| JSON_ENFORCED = "json_enforced" # constrained decoding guaranteed | |
| EXTRACTION = "extraction" # GLiNER2/NuExtract-style native | |
| @dataclass | |
| class EndpointConfig: | |
| capability: EndpointCapability | |
| call_fn: Callable[[str], str] | |
| # call_fn signature: takes a prompt string, returns a response string | |
| # caller is responsible for auth, model selection, retries | |
| ``` | |
| The `call_fn` injection means the library is agnostic about whether the caller is hitting Anthropic, OpenAI, a local Ollama instance, or Fastino's GLiNER2 API. The library just hands it a string and gets a string back. | |
| --- | |
| ## Pipeline (`pipeline.py`) | |
| ```python | |
| @dataclass | |
| class AnnotationResult: | |
| xml: str # annotated XML string | |
| fuzzy_spans: list[ResolvedSpan] # spans flagged for human review | |
| def annotate( | |
| text: str, # may contain existing XML tags | |
| schema: TEISchema, # subset of TEI elements in scope | |
| endpoint: EndpointConfig, # injected inference dependency | |
| gliner_model: str | None = "numind/NuNER_Zero", # None disables GLiNER pass | |
| chunk_size: int = 1500, # chars | |
| chunk_overlap: int = 200, | |
| ) -> AnnotationResult: | |
| ``` | |
| ### Execution Flow | |
| ``` | |
| 1. SETUP | |
| strip existing XML tags from text for processing, | |
| preserve them as a restoration map for final merge | |
| 2. GLINER PASS (skipped if gliner_model=None, endpoint is EXTRACTION, or text is short) | |
| map TEISchema elements → flat label list for GLiNER | |
| e.g. [("persName", "a person's name"), ("placeName", "a place name"), ...] | |
| chunk text if len(text) > chunk_size (with overlap) | |
| run gliner.predict_entities() on each chunk | |
| merge cross-chunk duplicates by span text + context overlap | |
| output: list[SpanDescriptor] with text + context + element + confidence | |
| (GLiNER is a pre-filter only; the LLM may reject, correct, or extend its candidates) | |
| 3. PROMPT ASSEMBLY | |
| select template based on EndpointCapability: | |
| TEXT_GENERATION: include JSON structure example + "output only JSON" instruction | |
| JSON_ENFORCED: minimal prompt, schema enforced externally | |
| EXTRACTION: pass schema directly in endpoint's native format, skip LLM prompt | |
| inject into prompt: | |
| - TEIElement descriptions + allowed attributes for in-scope elements | |
| - GLiNER pre-detected spans as candidates for the model to enrich/correct | |
| - source text chunk | |
| - instruction to emit one SpanDescriptor per occurrence, not per unique entity | |
| 4. INFERENCE | |
| call endpoint.call_fn(prompt) → raw response string | |
| 5. POSTPROCESSING (per chunk, then merged) | |
| a. Parse | |
| JSON_ENFORCED/EXTRACTION: parse directly | |
| TEXT_GENERATION: strip markdown fences, parse JSON, | |
| on failure: retry once with a correction prompt that includes | |
| the original (bad) response and the parse error message, | |
| so the model can self-correct rather than starting from scratch | |
| b. Resolve (resolver.py) | |
| for each SpanDescriptor: | |
| find context string in source → exact match preferred | |
| find text within context window | |
| assert source[start:end] == span.text → reject on mismatch | |
| fuzzy fallback (threshold 0.92) → flag for review | |
| c. Validate (validator.py) | |
| reject spans where text not in source | |
| check attributes against TEISchema allowed values | |
| check element is in schema scope | |
| d. Inject (injector.py) | |
| infer nesting from offset containment (child ⊂ parent by [start, end] bounds) | |
| check inferred nesting: children must be within parent bounds | |
| sort ResolvedSpans by start offset, handle nesting depth-first | |
| insert tags into a copy of the original source string | |
| restore previously existing XML tags from step 1 | |
| e. Final validation | |
| parse output as XML → reject malformed documents | |
| optionally validate against full TEI RelaxNG schema via lxml | |
| 6. RETURN | |
| AnnotationResult( | |
| xml=annotated_xml_string, | |
| fuzzy_spans=list_of_flagged_resolved_spans, | |
| ) | |
| ``` | |
| --- | |
| ## Key Design Constraints | |
| - The source text is **never modified by any model call**. All text in the output comes from the original input; models only contribute tag positions and attributes. | |
| - The **GLiNER pass is optional** (`gliner_model=None` disables it). It is most useful for long texts with `TEXT_GENERATION` endpoints; it is skipped automatically for `EXTRACTION` endpoints or short inputs. When enabled, GLiNER is a pre-filter only — the LLM may reject, correct, or extend its candidates. | |
| - **Span nesting is inferred from offsets**, never emitted by the model. `SpanDescriptor` is always flat; `ResolvedSpan.children` is populated by the injector from containment relationships. | |
| - `call_fn` has **no required signature beyond `(str) -> str`**, making it trivial to swap endpoints, add logging, or inject mock functions for testing. | |
| - Fuzzy-matched spans are **surfaced, not silently accepted** — `AnnotationResult.fuzzy_spans` provides a reviewable list alongside the XML. | |
| --- | |
| ## Testing Strategy | |
| ### Mocking philosophy | |
| **Always mock `call_fn` and the GLiNER detector in unit tests.** Do not use a real GLiNER model as a substitute for a remote LLM endpoint — GLiNER is a span-labelling model that cannot produce JSON responses; it cannot exercise the parse/resolve/inject pipeline. Using a real model also makes tests slow (~seconds per inference on CPU), non-deterministic across versions and hardware, and dependent on a 400MB+ download. | |
| The `call_fn: (str) -> str` design makes mocking trivial — a lambda returning a hardcoded JSON string is sufficient. No mock framework is needed. | |
| ### Test layers | |
| **Layer 1 — Unit tests** (always run, <1s total, fully mocked): | |
| ``` | |
| tests/ | |
| ├── test_chunker.py # chunker unit tests | |
| ├── test_resolver.py # resolver unit tests | |
| ├── test_validator.py # validator unit tests | |
| ├── test_injector.py # injector unit tests | |
| ├── test_builder.py # prompt builder unit tests | |
| ├── test_parser.py # JSON parse + retry unit tests | |
| └── test_pipeline.py # full pipeline smoke test (mocked call_fn + GLiNER) | |
| ``` | |
| **Layer 2 — Integration tests** (opt-in, gated by `pytest -m integration`): | |
| ``` | |
| tests/integration/ | |
| ├── test_gliner_detector.py # real GLiNER model, real HuggingFace download | |
| └── test_pipeline_e2e.py # full annotate() with real GLiNER + mocked call_fn | |
| ``` | |
| Integration tests are excluded from CI by default via `pyproject.toml`: | |
| ```toml | |
| [tool.pytest.ini_options] | |
| addopts = "-m 'not integration'" | |
| ``` | |
| ### TDD cycle (red → green → refactor) | |
| Each module is written test-first. Write a failing test, implement the minimum code to pass, refactor. | |
| ### Key test cases per module | |
| **`chunker.py`** | |
| - Short text below `chunk_size` → single chunk, offset 0 | |
| - Long text → multiple chunks with correct `start_offset` per chunk | |
| - Span exactly at a chunk boundary → appears in both chunks with correct global offset | |
| - Input with existing XML tags → chunk boundaries never split a tag | |
| **`resolver.py`** | |
| - Exact context match → `ResolvedSpan` with `fuzzy_match=False` | |
| - Context not found in source → span rejected | |
| - `source[start:end] != span.text` → span rejected | |
| - Context found but span text not within it → span rejected | |
| - Context found, span text found, score < 0.92 → span rejected | |
| - Context found, span text found, 0.92 ≤ score < 1.0 → `fuzzy_match=True` | |
| - Multiple occurrences of context → first match used, or rejection if ambiguous | |
| **`validator.py`** | |
| - Element not in schema → span rejected | |
| - Attribute not in schema → span rejected | |
| - Attribute value not in `allowed_values` → span rejected | |
| - Valid span → passes through unchanged | |
| **`injector.py`** | |
| - Two non-overlapping spans → both tags inserted correctly | |
| - Span B offset-contained in span A → B is child of A in output | |
| - Overlapping (non-nesting) spans → reject or flatten with warning | |
| - Restored XML tags from step 1 do not conflict with injected tags | |
| **`builder.py`** | |
| - `TEXT_GENERATION` capability → prompt contains JSON example and "output only JSON" instruction | |
| - `JSON_ENFORCED` capability → prompt is minimal (no JSON scaffolding) | |
| - GLiNER candidates present → candidates appear in prompt | |
| - GLiNER candidates absent (pass skipped) → prompt has no candidate section | |
| **`parser.py` (parse + retry logic)** | |
| - Valid JSON response → parsed to `list[SpanDescriptor]` without retry | |
| - Markdown-fenced JSON → fences stripped, parsed correctly | |
| - Invalid JSON on first attempt → retry triggered with correction prompt that includes original bad response + parse error message | |
| - Invalid JSON on second attempt → exception raised, chunk skipped | |
| **`pipeline.py` (smoke test)** | |
| ```python | |
| def mock_call_fn(prompt: str) -> str: | |
| return json.dumps([ | |
| {"element": "persName", "text": "John Smith", | |
| "context": "...said John Smith yesterday...", "attrs": {}} | |
| ]) | |
| def test_annotate_smoke(): | |
| schema = TEISchema(elements=[ | |
| TEIElement(tag="persName", description="a person's name", | |
| allowed_children=[], attributes=[]) | |
| ]) | |
| endpoint = EndpointConfig( | |
| capability=EndpointCapability.JSON_ENFORCED, | |
| call_fn=mock_call_fn, | |
| ) | |
| result = annotate( | |
| text="He said John Smith yesterday.", | |
| schema=schema, | |
| endpoint=endpoint, | |
| gliner_model=None, # disable GLiNER in unit tests | |
| ) | |
| assert "persName" in result.xml | |
| assert "John Smith" in result.xml | |
| assert result.xml.count("John Smith") == 1 # text not duplicated | |
| ``` | |
| --- | |
| ## Implementation Status | |
| **Completed 2026-02-28** — full implementation per the plan above. | |
| ### What was built | |
| All modules in the package structure were implemented: | |
| | File | Notes | | |
| | --- | --- | | |
| | `tei_annotator/models/schema.py` | `TEIAttribute`, `TEIElement`, `TEISchema` dataclasses | | |
| | `tei_annotator/models/spans.py` | `SpanDescriptor`, `ResolvedSpan` dataclasses | | |
| | `tei_annotator/inference/endpoint.py` | `EndpointCapability` enum, `EndpointConfig` dataclass | | |
| | `tei_annotator/chunking/chunker.py` | `chunk_text()` — overlap chunker, XML-safe boundaries | | |
| | `tei_annotator/detection/gliner_detector.py` | `detect_spans()` — optional, raises `ImportError` if `[gliner]` extra not installed | | |
| | `tei_annotator/prompting/builder.py` | `build_prompt()` + `make_correction_prompt()` | | |
| | `tei_annotator/prompting/templates/text_gen.jinja2` | Verbose prompt with JSON example, "output only JSON" instruction | | |
| | `tei_annotator/prompting/templates/json_enforced.jinja2` | Minimal prompt for constrained-decoding endpoints | | |
| | `tei_annotator/postprocessing/parser.py` | `parse_response()` — fence stripping, one-shot self-correction retry | | |
| | `tei_annotator/postprocessing/resolver.py` | `resolve_spans()` — context-anchor → char offset, rapidfuzz fuzzy fallback at threshold 0.92 | | |
| | `tei_annotator/postprocessing/validator.py` | `validate_spans()` — element, attribute name, allowed-value checks | | |
| | `tei_annotator/postprocessing/injector.py` | `inject_xml()` — stack-based nesting tree, recursive tag insertion | | |
| | `tei_annotator/pipeline.py` | `annotate()` — full orchestration, tag strip/restore, deduplication across chunks, lxml final validation | | |
| ### Dependencies added | |
| Runtime: `jinja2`, `lxml`, `rapidfuzz`. Optional extra `[gliner]` for GLiNER support. Dev: `pytest`, `pytest-cov`. | |
| ### Tests | |
| - **63 unit tests** (Layer 1) — fully mocked, run in < 0.1 s via `uv run pytest` | |
| - **9 integration tests** (Layer 2, no GLiNER) — complex resolver/injector/pipeline scenarios, run via `uv run pytest --override-ini="addopts=" -m integration tests/integration/test_pipeline_e2e.py -k "not real_gliner"` | |
| - **1 GLiNER integration test** — requires `[gliner]` extra and HuggingFace model download | |
| ### Smoke script | |
| `scripts/smoke_test_llm.py` — end-to-end test with real LLM calls (no GLiNER). Verified against: | |
| - **Google Gemini 2.0 Flash** (`GEMINI_API_KEY` from `.env`) | |
| - **KISSKI `llama-3.3-70b-instruct`** (`KISSKI_API_KEY` from `.env`, OpenAI-compatible API at `https://chat-ai.academiccloud.de/v1`) | |
| Run with `uv run scripts/smoke_test_llm.py`. | |
| ### Key implementation notes | |
| - The `_strip_existing_tags` / `_restore_existing_tags` pair in `pipeline.py` preserves original markup by tracking plain-text offsets of each stripped tag and re-inserting them after annotation. | |
| - `_build_nesting_tree` in `injector.py` uses a sort-by-(start-asc, length-desc) + stack algorithm; partial overlaps are dropped with a `warnings.warn`. | |
| - The resolver does an exact `str.find` first; fuzzy search (sliding-window rapidfuzz) is only attempted if exact fails and rapidfuzz is installed. | |
| - `parse_response` passes `call_fn` and `make_correction_prompt` only for `TEXT_GENERATION` endpoints; `JSON_ENFORCED` and `EXTRACTION` never retry. | |
| --- | |
| ## RNG Schema Parsing (`tei_annotator/tei.py`) | |
| **Added 2026-02-28** — `create_schema()` parses a RELAX NG schema file and produces a `TEISchema` for use with `annotate()`. | |
| ### Function signature | |
| ```python | |
| def create_schema( | |
| schema_path: str | Path, | |
| element: str = "text", | |
| depth: int = 1, | |
| ) -> TEISchema | |
| ``` | |
| - `schema_path` — path to a `.rng` file (e.g. `schema/tei-bib.rng`). | |
| - `element` — name of the root TEI element to start from (default: `"text"`). | |
| - `depth` — how many levels of descendant elements to include. `depth=1` adds the root element **and** its direct children; `depth=2` adds grandchildren too; `depth=0` adds only the root itself. | |
| - Raises `ValueError` if the element is not found in the schema. | |
| ### How it works | |
| 1. Parses the `.rng` with lxml, builds a `{define_name: element_node}` lookup table. | |
| 2. Builds a reverse map from TEI element names (e.g. `"persName"`) to their RNG define names (e.g. `"tbibpersName"`). | |
| 3. BFS from the requested element, collecting `TEIElement` entries level by level up to `depth`. | |
| 4. For each element: | |
| - **description** — extracted from the first `<a:documentation>` child inside the RNG `<element>` node. | |
| - **`allowed_children`** — content-model `<ref>` nodes are expanded recursively through macro/model groups; attribute-group refs (names containing `"att."`) are skipped; element-bearing defines are short-circuited (just record the element name, don't recurse — correctly handles self-referential elements like `idno`). | |
| - **`attributes`** — attribute-group refs are followed recursively; inline `<attribute>` elements are also collected; `required` is inferred from the immediate parent (`<optional>` / `<zeroOrMore>` → False, otherwise True); `allowed_values` comes from `<choice><value>` enumeration inside the attribute definition. | |
| 5. Deduplicates children and attributes (preserving order) before constructing each `TEIElement`. | |
| ### Tests (`tests/test_tei.py`) | |
| 15 unit tests — run in < 0.05 s, no network, no model downloads: | |
| - `idno` (leaf-like, self-referential, enumerated `type` values: `ISBN`, `ISSN`, `DOI`, …) | |
| - `biblStruct` (explicit named children: `analytic`, `monogr`, `series`, `relatedItem`, `citedRange`; plus model-group expansion: `note`, `ptr`) | |
| - `depth=0` vs `depth=1` behaviour | |
| - Duplicate element / attribute detection | |
| - Unknown element raises `ValueError` | |
| Total unit tests after this addition: **78** (63 original + 15 new). | |
| --- | |
| ## Evaluation Module (`tei_annotator/evaluation/`) | |
| **Added 2026-02-28** — compares annotator output against a gold-standard TEI XML file to compute precision, recall, and F1 score. | |
| ### Package structure | |
| ``` | |
| tei_annotator/evaluation/ | |
| ├── __init__.py # public API exports | |
| ├── extractor.py # EvaluationSpan, extract_spans(), spans_from_xml_string() | |
| ├── metrics.py # MatchMode, SpanMatch, ElementMetrics, EvaluationResult, | |
| │ # match_spans(), compute_metrics(), aggregate() | |
| └── evaluator.py # evaluate_bibl(), evaluate_file() | |
| ``` | |
| ### Algorithm | |
| For each gold-standard element (e.g. `<bibl>`): | |
| 1. **Extract gold spans** — walk the element tree, record `(tag, start, end, text)` for every descendant element using absolute char offsets in the element's plain text (`"".join(element.itertext())`). | |
| 2. **Strip tags** — the same plain text is passed to `annotate()`. | |
| 3. **Annotate** — run the full pipeline with the injected `EndpointConfig`. | |
| 4. **Extract predicted spans** — parse `result.xml` as `<_root>…</_root>`, then run the same span extractor. | |
| 5. **Match** — greedily pair (gold, pred) spans by score (highest first); each span matched at most once. | |
| 6. **Compute metrics** — count TP/FP/FN per element type; derive precision, recall, F1. | |
| ### Match modes (`MatchMode`) | |
| | Mode | A match if… | | |
| | --- | --- | | |
| | `TEXT` (default) | same element tag + normalised text content | | |
| | `EXACT` | same element tag + identical `(start, end)` offsets | | |
| | `OVERLAP` | same element tag + IoU ≥ `overlap_threshold` (default 0.5) | | |
| ### Key entry points | |
| ```python | |
| # Single element evaluation (lxml element) | |
| result = evaluate_bibl(gold_element, schema, endpoint, gliner_model=None) | |
| print(result.report()) | |
| # Full file evaluation | |
| per_record, overall = evaluate_file( | |
| "tests/fixtures/blbl-examples.tei.xml", | |
| schema=schema, | |
| endpoint=endpoint, | |
| max_items=10, # optional — first N bibls only | |
| ) | |
| print(overall.report()) | |
| ``` | |
| `EvaluationResult` exposes: | |
| - `micro_precision / micro_recall / micro_f1` — aggregate counts, then compute rates | |
| - `macro_precision / macro_recall / macro_f1` — average per-element rates | |
| - `per_element: dict[str, ElementMetrics]` — per-element breakdown | |
| - `matched / unmatched_gold / unmatched_pred` — full span lists for inspection | |
| - `report()` — human-readable summary string | |
| ### Evaluation Tests | |
| 39 new unit tests in `tests/test_evaluation.py` — all mocked, run in < 0.1 s. | |
| Total unit tests: **117** (78 + 39). | |