"""Human-annotated multilingual inference records for the graph corpus.""" from __future__ import annotations import hashlib import json import re import zipfile from collections.abc import Iterator, Mapping from pathlib import Path from graph.graph import EdgeRecord, GraphStore from graph.real.corpus import SourceMetadata from graph.synthetic.generator import ReasoningExample from graph.types import GraphPath, GraphQuery, PathStep, ReasoningStatus LABELS = ("entailment", "neutral", "contradiction") def _stable_id(value: str) -> str: return hashlib.blake2b(value.encode("utf-8"), digest_size=8).hexdigest() def iter_xnli_rows(archive: str | Path) -> Iterator[dict[str, str]]: """Yield normalized XNLI JSONL rows from the official ZIP artifact.""" with zipfile.ZipFile(archive) as handle: names = sorted( name for name in handle.namelist() if name.rsplit("/", 1)[-1] in {"xnli.dev.jsonl", "xnli.test.jsonl"} ) for name in names: split = Path(name).stem.removeprefix("xnli.") with handle.open(name) as binary: for line in binary: if not line.strip(): continue raw = json.loads(line) if not isinstance(raw, Mapping): continue label = str(raw.get("gold_label", "")).strip().lower() language = str(raw.get("language", "")).strip().lower() premise = " ".join( str(raw.get("premise", raw.get("sentence1", ""))).split() ) hypothesis = " ".join( str(raw.get("hypothesis", raw.get("sentence2", ""))).split() ) pair_id = str(raw.get("pairID", raw.get("pair_id", ""))) if ( label not in LABELS or not language or not premise or not hypothesis ): continue yield { "pair_id": pair_id or _stable_id(premise + "\n" + hypothesis), "split": split, "language": language, "premise": premise, "hypothesis": hypothesis, "label": label, } def xnli_to_reasoning_example( row: Mapping[str, str], source: SourceMetadata ) -> ReasoningExample: key = f"{row['language']}:{row['pair_id']}" suffix = _stable_id(key) premise_id = f"xnli:{suffix}:premise" hypothesis_id = f"xnli:{suffix}:hypothesis" label_ids = {label: f"xnli:label:{label}" for label in LABELS} nodes = { premise_id: "premise", hypothesis_id: "hypothesis", **{node_id: "inference_label" for node_id in label_ids.values()}, } attributes = { premise_id: { "text": row["premise"], "language": row["language"], "human_annotation": "XNLI", }, hypothesis_id: {"text": row["hypothesis"], "language": row["language"]}, **{ node_id: {"label": label, "candidate": "true"} for label, node_id in label_ids.items() }, } edges = [ EdgeRecord( premise_id, "has_hypothesis", hypothesis_id, "premise", "hypothesis", "a premise-to-hypothesis connection", {"evidence_kind": "human_annotated_multilingual_inference"}, ), EdgeRecord( hypothesis_id, "has_inference_label", label_ids[row["label"]], "hypothesis", "inference_label", "an annotated inference-label connection", {"evidence_kind": "human_annotated_multilingual_inference"}, ), ] graph = GraphStore( nodes, edges, graph_version=f"{source.source_id}@{source.version}", node_attributes=attributes, relation_attributes={ "has_hypothesis": {"functional": "true"}, "has_inference_label": {"functional": "true"}, }, ) path = GraphPath( ( PathStep( premise_id, "has_hypothesis", hypothesis_id, "premise", "hypothesis", edges[0].attributes, ), PathStep( hypothesis_id, "has_inference_label", label_ids[row["label"]], "hypothesis", "inference_label", edges[1].attributes, ), ) ) query = GraphQuery( text=( f"In {row['language']}, determine the inference class for this premise: " f"{row['premise']} Hypothesis: {row['hypothesis']}" ), start_entities=(premise_id,), start_entity_ids=(premise_id,), min_hops=2, max_hops=2, exact_hops=2, return_paths=True, verify=True, operation="find_entities", answer_mode="entity", ) seed = int.from_bytes( hashlib.blake2b(key.encode(), digest_size=8).digest(), "little" ) return ReasoningExample( example_id=f"xnli-{row['language']}-{row['pair_id']}", graph=graph.to_example(), query=query, answers=(label_ids[row["label"]],), negative_answers=tuple( label_ids[label] for label in LABELS if label != row["label"] ), supporting_paths=(path,), conflicting_paths=(), status=ReasoningStatus.SUPPORTED, difficulty=2, seed=seed, generator_version="xnli-corpus-v1", task_type="multilingual_inference", intent_id=f"xnli:{row['language']}:{row['label']}", composition=("has_hypothesis", "has_inference_label"), stage_name=f"multilingual_xnli_{row['split']}", ) def _seed_language(filename: str) -> str | None: """Extract an NLLB language code from common Seed archive filenames.""" base = Path(filename).name stem = base.rsplit(".", 1)[0] if stem.lower() in {"readme", "license", "metadata", "languages", "seed"}: return None if re.fullmatch(r"[a-z]{3}(?:_[A-Za-z]+)?", stem): return stem for token in re.split(r"[._-]", stem): if re.fullmatch(r"[a-z]{3}(?:_[A-Za-z]+)?", token): return token return None def iter_nllb_seed_rows( archive: str | Path, *, max_records: int | None = None ) -> Iterator[dict[str, str]]: """Yield English-pivot aligned rows from the professionally translated seed set.""" with zipfile.ZipFile(archive) as handle: files: dict[str, list[str]] = {} for name in handle.namelist(): if name.endswith("/"): continue language = _seed_language(name) if language is None: continue with handle.open(name) as binary: rows = [ " ".join(line.decode("utf-8", "replace").split()) for line in binary ] rows = [row for row in rows if row] if rows: files[language] = rows if not files: raise ValueError("NLLB-Seed archive contains no language files") anchor = next( (language for language in files if language in {"eng", "eng_Latn"}), min(files), ) index = 0 for target_language in sorted(files): if target_language == anchor: continue for sentence_index, (source_text, target_text) in enumerate( zip(files[anchor], files[target_language], strict=False) ): yield { "pair_id": f"{sentence_index}", "source_language": anchor, "target_language": target_language, "source_text": source_text, "target_text": target_text, } index += 1 if max_records is not None and index >= max_records: return def nllb_seed_to_reasoning_example( row: Mapping[str, str], source: SourceMetadata ) -> ReasoningExample: """Represent one aligned NLLB-Seed pair as a graph traversal example.""" key = f"{row['source_language']}:{row['target_language']}:{row['pair_id']}" suffix = _stable_id(key) source_id = f"nllb:{suffix}:source" target_id = f"nllb:{suffix}:target" edge = EdgeRecord( source_id, "translation_of", target_id, "sentence", "sentence", "a translation-alignment connection", { "evidence_kind": "human_translated_multilingual", "source_language": row["source_language"], "target_language": row["target_language"], }, ) graph = GraphStore( {source_id: "sentence", target_id: "sentence"}, [edge], graph_version=f"{source.source_id}@{source.version}", node_attributes={ source_id: { "text": row["source_text"], "language": row["source_language"], }, target_id: { "text": row["target_text"], "language": row["target_language"], }, }, relation_attributes={"translation_of": {"functional": "true"}}, ) path = GraphPath( ( PathStep( source_id, "translation_of", target_id, "sentence", "sentence", edge.attributes, ), ) ) query = GraphQuery( text=( f"In {row['target_language']}, identify the sentence aligned with this " f"{row['source_language']} sentence: {row['source_text']}" ), start_entities=(source_id,), start_entity_ids=(source_id,), min_hops=1, max_hops=1, exact_hops=1, return_paths=True, verify=True, operation="find_entities", answer_mode="entity", ) return ReasoningExample( example_id=f"nllb-seed-{key}", graph=graph.to_example(), query=query, answers=(target_id,), negative_answers=(), supporting_paths=(path,), conflicting_paths=(), status=ReasoningStatus.SUPPORTED, difficulty=1, seed=int.from_bytes( hashlib.blake2b(key.encode(), digest_size=8).digest(), "little" ), generator_version="nllb-seed-corpus-v1", task_type="multilingual_alignment", intent_id=f"nllb_seed:{row['target_language']}", composition=("translation_of",), stage_name="multilingual_nllb_seed", )