189 lines
6.8 KiB
Python
189 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
||
"""Convert Deep Research source IDs into numeric citations for final output."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
SRC_CITE_RE = re.compile(r"\[((?:src_[A-Za-z0-9_-]+)(?:\s*,\s*src_[A-Za-z0-9_-]+)*)\]")
|
||
|
||
|
||
def load_sources(path: Path) -> dict[str, dict[str, Any]]:
|
||
sources: dict[str, dict[str, Any]] = {}
|
||
if not path.exists():
|
||
return sources
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
if not line.strip():
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
sid = obj.get("id") or obj.get("source_id")
|
||
if sid:
|
||
sources[str(sid)] = obj
|
||
return sources
|
||
|
||
|
||
def extract_ordered_source_ids(text: str) -> list[str]:
|
||
ordered: list[str] = []
|
||
seen: set[str] = set()
|
||
for match in SRC_CITE_RE.finditer(text):
|
||
for sid in [item.strip() for item in match.group(1).split(",")]:
|
||
if sid and sid not in seen:
|
||
seen.add(sid)
|
||
ordered.append(sid)
|
||
return ordered
|
||
|
||
|
||
def _canonical_source_key(sid: str, source: dict[str, Any] | None) -> str:
|
||
"""Return a stable de-duplication key for a source record.
|
||
|
||
Phase 2 often creates chapter-local source IDs for the same local PDF or
|
||
official guideline. Final references should cite the underlying source
|
||
once, while citation_map.json keeps the full src_id traceability.
|
||
"""
|
||
if not source:
|
||
return f"missing:{sid}"
|
||
title = re.sub(r"\s+", " ", str(source.get("title") or source.get("name") or sid)).strip().lower()
|
||
title = title.removesuffix(" ocr").removesuffix(".ocr").strip()
|
||
doi = str(source.get("doi") or "").strip().lower()
|
||
if doi:
|
||
return f"doi:{doi}"
|
||
path = str(source.get("path") or "").strip()
|
||
url = str(source.get("url") or "").strip()
|
||
if title and ("phase0/extracted/" in path or "phase0/extracted/" in url):
|
||
return f"local-material:{title}"
|
||
for field in ("url", "path"):
|
||
value = str(source.get(field) or "").strip()
|
||
if value:
|
||
return f"{field}:{value.rstrip('/').lower()}"
|
||
return f"title:{title or sid}"
|
||
|
||
|
||
def build_numeric_mapping(
|
||
ordered_ids: list[str],
|
||
sources: dict[str, dict[str, Any]],
|
||
) -> tuple[dict[str, int], list[dict[str, Any]]]:
|
||
mapping: dict[str, int] = {}
|
||
records: list[dict[str, Any]] = []
|
||
seen_keys: dict[str, int] = {}
|
||
record_by_number: dict[int, dict[str, Any]] = {}
|
||
for sid in ordered_ids:
|
||
source = sources.get(sid)
|
||
key = _canonical_source_key(sid, source)
|
||
if key in seen_keys:
|
||
number = seen_keys[key]
|
||
mapping[sid] = number
|
||
record_by_number[number].setdefault("source_ids", []).append(sid)
|
||
continue
|
||
number = len(records) + 1
|
||
seen_keys[key] = number
|
||
mapping[sid] = number
|
||
record = {
|
||
"number": number,
|
||
"source_id": sid,
|
||
"source_ids": [sid],
|
||
"source": source or {},
|
||
"dedupe_key": key,
|
||
}
|
||
records.append(record)
|
||
record_by_number[number] = record
|
||
return mapping, records
|
||
|
||
|
||
def format_reference(number: int, sid: str, source: dict[str, Any] | None) -> str:
|
||
if not source:
|
||
return f"{number}. {sid}. （sources.jsonl 未找到该来源）"
|
||
authors = ", ".join(source.get("authors", [])) if source.get("authors") else ""
|
||
year = source.get("year") or source.get("date") or ""
|
||
title = source.get("title") or source.get("name") or sid
|
||
title = re.sub(r"(?i)(?:\s+OCR|\.ocr)$", "", str(title)).strip()
|
||
publisher = source.get("publisher") or source.get("venue") or source.get("source") or ""
|
||
url = source.get("url") or source.get("path") or ""
|
||
parts = [f"{number}. "]
|
||
if authors:
|
||
parts.append(f"{authors}. ")
|
||
if year:
|
||
parts.append(f"({year}). ")
|
||
parts.append(str(title))
|
||
if publisher:
|
||
parts.append(f". {publisher}")
|
||
if url:
|
||
parts.append(f". {url}")
|
||
return "".join(parts)
|
||
|
||
|
||
def convert_citations(text: str, mapping: dict[str, int]) -> str:
|
||
def repl(match: re.Match[str]) -> str:
|
||
ids = [item.strip() for item in match.group(1).split(",") if item.strip()]
|
||
nums: list[str] = []
|
||
seen: set[int] = set()
|
||
for sid in ids:
|
||
if sid not in mapping:
|
||
continue
|
||
number = mapping[sid]
|
||
if number in seen:
|
||
continue
|
||
seen.add(number)
|
||
nums.append(str(number))
|
||
return "<sup>[" + ", ".join(nums) + "]</sup>" if nums else match.group(0)
|
||
|
||
return SRC_CITE_RE.sub(repl, text)
|
||
|
||
|
||
def strip_existing_reference_section(text: str) -> str:
|
||
pattern = re.compile(r"\n##\s*(?:参考文献|参考来源清单|References)\s*\n.*\Z", re.S)
|
||
return pattern.sub("", text).rstrip() + "\n"
|
||
|
||
|
||
def number_citations(
|
||
*,
|
||
text: str,
|
||
sources: dict[str, dict[str, Any]],
|
||
) -> tuple[str, list[dict[str, Any]]]:
|
||
ordered_ids = extract_ordered_source_ids(text)
|
||
mapping, records = build_numeric_mapping(ordered_ids, sources)
|
||
body = convert_citations(strip_existing_reference_section(text), mapping).rstrip()
|
||
ref_lines = ["", "## 参考来源清单", ""]
|
||
for record in records:
|
||
ref_lines.append(format_reference(record["number"], record["source_id"], record["source"]))
|
||
return body + "\n" + "\n".join(ref_lines).rstrip() + "\n", records
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(description="Convert [src_xxx] citations to numeric citations")
|
||
parser.add_argument("project", help="Project directory")
|
||
parser.add_argument("--input", default="phase4/final_zh.md")
|
||
parser.add_argument("--output", default="phase4/final_zh_numbered.md")
|
||
parser.add_argument("--sources", default="phase2/sources.jsonl")
|
||
parser.add_argument("--map", default="phase4/citation_map.json")
|
||
args = parser.parse_args()
|
||
|
||
project = Path(args.project)
|
||
src_path = project / args.input
|
||
out_path = project / args.output
|
||
sources_path = project / args.sources
|
||
map_path = project / args.map
|
||
if not src_path.exists():
|
||
raise SystemExit(f"input not found: {src_path}")
|
||
sources = load_sources(sources_path)
|
||
numbered, records = number_citations(text=src_path.read_text(encoding="utf-8"), sources=sources)
|
||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
out_path.write_text(numbered, encoding="utf-8")
|
||
map_path.parent.mkdir(parents=True, exist_ok=True)
|
||
map_path.write_text(json.dumps(records, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||
print(f"Wrote: {out_path.relative_to(project)}")
|
||
print(f"Wrote: {map_path.relative_to(project)}")
|
||
print(f"Citations: {len(records)}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|