citation-extractor-fast

An experimental legal-citation extractor built by full-weight fine-tuning NVIDIA Nemotron 3 Nano 4B BF16 for one epoch. It extracts citation text, citation kinds and source-offset predictions from legal documents.

Input/output contract

Supply the source text itself as the user message, preserving its original characters, whitespace and line endings. Use the trained system prompt shipped in SYSTEM_PROMPT.txt, SHA-256:

8e2451afcf15498b0aab55b490fd0536342fcf097ce5dc3abd99958556de4f85

Load the prompt file byte-for-byte. Apply the shipped chat template with add_generation_prompt=True, enable_thinking=False, then tokenize with add_special_tokens=False. Load the prompt, tokenizer and weights from the same immutable revision.

For a UTF-8 file, raw = Path("document.txt").read_bytes().decode("utf-8") preserves the source characters and line endings. Pass raw as the user content.

Intended output: a JSON array of {quote, start, end, kind} objects. Offsets are zero-based Unicode character/codepoint indices in the original text, with exclusive end. Repeated citation occurrences are separate records. [] is the model's prediction of an empty citation list. kind is a learned teacher label; taxonomy aliases, target normalization, destination linking and legal validity are separate downstream tasks. Use quote to identify the extracted citation text. Applications that need source highlighting can locate each quote in the original text; repeated matches require occurrence disambiguation. The generated start and end fields are unreliable, so check raw[start:end] == quote when using model-provided offsets. Keep derived alignments distinguishable from model predictions. Complete JSON arrays and explicitly recorded generation failures support reliable downstream processing.

Seen training example: 27 U.S.C. § 121

This complete section was among the 8,015 examples actually used for fine-tuning. Its recorded silver training target below illustrates the input/output format. Evaluation results appear in the following section. Download the exact input, output, or complete message record and provenance.

Input — user message (the system message is SYSTEM_PROMPT.txt):

§ 121. State statutes as operative on termination of transportation; original packages
All fermented, distilled, or other intoxicating liquors or liquids transported into any State or Territory or remaining therein for use, consumption, sale, or storage therein, shall upon arrival in such State or Territory be subject to the operation and effect of the laws of such State or Territory enacted in the exercise of its police powers, to the same extent and in the same manner as though such liquids or liquors had been produced in such State or Territory, and shall not be exempt therefrom by reason of being introduced therein in original packages or otherwise.
(Aug. 8, 1890, ch. 728, 26 Stat. 313.)

Statutory Notes and Related Subsidiaries
Short TitleAct Aug. 8, 1890, is popularly known as the “Wilson Act” or the “Original Packages Act”.

Output — recorded training target:

[
  {"quote": "Aug. 8, 1890, ch. 728, 26 Stat. 313", "start": 663, "end": 698, "kind": "statutes_at_large"},
  {"quote": "Act Aug. 8, 1890", "start": 754, "end": 770, "kind": "act"},
  {"quote": "Wilson Act", "start": 799, "end": 809, "kind": "act"},
  {"quote": "Original Packages Act", "start": 819, "end": 840, "kind": "act"}
]

The downloadable input retains the original Unicode spacing and all trailing line endings. Each displayed quote matches its source slice at the recorded offsets. Training-example membership and source SHA-256 are included in the message record.

Evaluation: federal silver-reference eval200

All 200 frozen documents remain in every denominator, including failures, 31 output caps and 19 per-document timeouts. The evaluation excludes training IDs, source families and exact-text duplicates, and samples acts, bills, proposed rules, regulations and statutes. These are teacher-consensus silver references. The primary results below measure citation quote recognition: an extracted citation earns credit when its exact quote matches a reference, irrespective of the model's character-offset arithmetic. The metric counts repeated quote occurrences as a multiset. This post-hoc measure evaluates quote recognition against the silver references.

Baseline / GPT‑6.1 / fine-tuned comparison by document type

Exact quote recognition F1, ignoring generated offsets. All attempted documents, including caps, timeouts and invalid outputs, remain in the original evaluation denominator. All three arms used the same frozen 200-document cohort.

Document type	Documents	Untrained Nemotron	GPT‑6.1 Sol	Fine-tuned original	Fine-tuned + timeout retries
Acts	37	3.32%	23.21%	55.76%	61.88%
Bills	44	0.73%	13.19%	49.01%	49.01%
Proposed rules	37	0.88%	31.43%	32.58%	32.58%
Regulations	44	0.59%	30.01%	38.08%	34.15%
Statutes	38	1.33%	17.57%	43.64%	43.64%
Overall (micro)	200	1.31%	23.18%	44.23%	44.55%

The full document-type comparison also includes precision, recall, kind agreement and failures for every arm. GPT‑6.1 used the frozen comparison prompt and API protocol. The table compares the three models under that matched prompt.

Fine-tuned model summary — same 200 documents

Metric	Original	With timeout retries
Exact quote recognition precision, ignoring offsets	72.63%	65.07%
Exact quote recognition recall, ignoring offsets	31.79%	33.87%
Exact quote recognition F1, ignoring offsets	44.23%	44.55%
Valid JSON documents	150/200	155/200
Output-cap failures	31/200	45/200
Timeouts	19/200	0/200

Larger-budget follow-up for the 19 timeouts

The retry column uses 181 original predictions plus new outcomes for the 19 original timeout cases, retaining the full 200-document denominator. The first 5090 retry increased the per-document deadline from 600 to 1,800 seconds. Four EOS-completed outputs were retained; the remaining 15 assignments then received 32,768 output tokens, a 65,536-token context budget and the same 1,800-second deadline. Thinking-disabled greedy decoding and model weights were unchanged. The extra budget applies to citation output.

The completed 19-case follow-up produced five complete JSON arrays and 14 capped outputs. Its quote precision/recall/F1 were 25.08% / 13.69% / 17.71%. Across the full 200 documents, F1 increased by 0.32 percentage points. Repeated citation objects and drifting offsets persisted in the longer generations.

The baseline and GPT columns retain their original predictions. GPT‑6.1 Sol's recorded requests used each document's original max_tokens allowance, up to 16,384, with low reasoning effort. The last column measures an explicitly failure-selected, larger-budget retry procedure. Full predictions, protocols, all-attempt accounting and reproducible scores are included in the dataset supplement.

Untrained Nemotron quote-only F1 was 1.31%; the GPT comparison arm was 23.18%. The fine-tune improves quote recognition but remains unreliable for exact offsets and complete extraction. Comparison runtimes span Spark/RTX 5090/API hardware, so latency reflects those different environments. Full aggregate and per-document-type reports are included under evaluation/. The 200-document cohort uses complete-prompt-plus-reference fit at 32,768 tokens and consists of federal documents.

Diagnostic: generated character offsets

The full report also retains exact-span-plus-kind scoring (0.32% F1 for this model). That diagnostic requires the quote, both numerical offsets and kind to match simultaneously. It is useful specifically for applications that consume the model's generated offsets directly. It is a secondary diagnostic: a correctly extracted citation with a miscounted character position fails it. The primary metric is citation quote recognition. Source alignment is a separate downstream operation. The original strict document/schema counts remain available in evaluation/comparison.json.

Training and weight provenance

• Base: nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16, revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f.
• Full language-weight fine-tuning with native BF16 parameters, architecture-required dtype exceptions and FP32 Adafactor optimizer state.
• 8,015 federal common-fit training examples; one epoch, 2,004 optimizer updates, microbatch one, accumulation four, learning rate 5e-6, 101 warmup updates, cosine schedule to 0.1 of initial LR, seed 102301.
• Common-fit set: 8,015 of 8,129 candidate training sources. The 114 union fit exclusions are recorded in the original audit; admitted inputs are complete.
• Training labels are reconciled teacher-consensus silver labels. Literal repeated occurrences were projected deterministically. Coverage omissions and teacher mistakes may remain.
• Final checkpoint: step 2,004, with complete export/reload equality verified for all 263 parameters. Every published weight shard must match that verified export. The weight-integrity equality report is included under provenance/.
• Export metadata reports the base model's larger positional limit; this fine-tune and evaluation only qualified complete sequences at 32,768 tokens. Use 32,768 tokens as the tested complete-sequence budget for this release.

The release includes SHA256SUMS.json and a training/provenance summary. Model weights and tokenizer bytes are unchanged from the evaluated export. The release generation_config.json additionally recognizes both native EOS (2) and trained chat EOS (11), matching the completed evaluation's stopping rule. This metadata correction preserves the weights and trained prompt. The original export generation config is retained under provenance/.

Training data and replication

The training and evaluation dataset contains the exact 8,015 training examples, 200 held-out evaluation examples, original predictions from all three comparison models, frozen training/scoring code, package versions, recipe and split manifests. Its verify_release.py checks artifact hashes, source offsets, split separation and reproduces every published evaluation score group from the supplied outputs. The dataset card provides the commands for loading data, verifying results and rerunning training from the pinned base model.

Usage

Use native Transformers 5.15.1, PyTorch and safetensors with BF16-capable hardware. Native Mamba/causal-convolution kernels strongly affect performance; upstream-supported CUDA packages or compatible pinned kernel builds may be needed for practical speed. Load the native Transformers implementation with trust_remote_code=False. The exact tested runtime versions and export checks are in the included qualification report. Install kernels appropriate to your CUDA environment before running the example.

import hashlib
from pathlib import Path
import torch
from huggingface_hub import HfApi, hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "thurstonpowers/citation-extractor-fast"
revision = HfApi().model_info(repo).sha  # pin this SHA for reproducible use
prompt = Path(hf_hub_download(repo, "SYSTEM_PROMPT.txt", revision=revision)).read_text()
assert hashlib.sha256(prompt.encode()).hexdigest() == "8e2451afcf15498b0aab55b490fd0536342fcf097ce5dc3abd99958556de4f85"
raw = Path(hf_hub_download(repo, "examples/seen-section-input.txt", revision=revision)).read_bytes().decode("utf-8")
tok = AutoTokenizer.from_pretrained(repo, revision=revision, trust_remote_code=False)
model = AutoModelForCausalLM.from_pretrained(
    repo, revision=revision, dtype=torch.bfloat16, trust_remote_code=False
).to("cuda").eval()
rendered = tok.apply_chat_template(
    [{"role": "system", "content": prompt}, {"role": "user", "content": raw}],
    tokenize=False, add_generation_prompt=True, enable_thinking=False,
)
inputs = tok(rendered, add_special_tokens=False, truncation=False, return_tensors="pt").to("cuda")
max_new_tokens = 16384
assert inputs.input_ids.shape[1] + max_new_tokens <= 32768, "Use a complete input and output budget within 32768 tokens"
eos = sorted({2, tok.eos_token_id})
with torch.inference_mode():
    output = model.generate(
        **inputs, do_sample=False, use_cache=True, max_new_tokens=max_new_tokens,
        max_time=600, eos_token_id=eos, pad_token_id=tok.pad_token_id,
    )
generated = output[0, inputs.input_ids.shape[1]:].tolist()
assert generated and generated[-1] in eos, "Record this generation as incomplete"
text = tok.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)
print(text)  # parse JSON and validate every offset/quote before use

This runs the seen training-section example above; generated output can differ from its recorded training target. inference.py includes strict output/source validation and preserves raw output on failure. Its 600-second generation setting is a best-effort Transformers limit. Use an external watchdog for a hard process deadline. The helper preserves complete input and raw model output, validates source offsets, and records incomplete generations as failures.

Limitations and intended use

Research, error analysis and human-reviewed citation candidate generation. Missing citations, wrong kinds, wrong offsets, repetition, malformed JSON, caps and timeouts remain material limitations. Use human review and independent authority resolution for consequential legal applications. Checking whether an authority exists, applies and remains current is a separate task from recognizing its citation text. Reported results cover the federal evaluation described above; other source distributions require their own evaluation.

License and attribution

Derived from NVIDIA's base checkpoint under the NVIDIA Nemotron Open Model License. A copy of the governing license is included as LICENSE.pdf, captured from the official NVIDIA license link. Original model-card attribution is retained under provenance/BASE_MODEL_CARD.md; see NOTICE for modification attribution.