codebook / potato /export /conll_u_exporter.py
davidjurgens's picture
Deploy: Potato — Codebook Annotation
aceb1b2 verified
Raw
History Blame Contribute Delete
6.16 kB
"""
CoNLL-U Exporter
Exports span annotations to CoNLL-U format:
- 10 columns: ID FORM LEMMA UPOS XPOS FEATS HEAD DEPREL DEPS MISC
- NER annotations placed in MISC column as SpaceAfter/NER features
- Blank lines between sentences
- Comment lines with sent_id and text
"""
import os
import logging
from typing import Optional, Tuple
from .base import BaseExporter, ExportContext, ExportResult
from .nlp_utils import tokenize_text, char_spans_to_bio_tags, group_sentences
logger = logging.getLogger(__name__)
class CoNLLUExporter(BaseExporter):
format_name = "conll_u"
description = "CoNLL-U format (Universal Dependencies compatible, NER in MISC)"
file_extensions = [".conllu"]
def can_export(self, context: ExportContext) -> Tuple[bool, str]:
has_span_schema = any(
s.get("annotation_type") == "span"
for s in context.schemas
)
if not has_span_schema:
return False, "No span annotation schema found in config"
return True, ""
def export(self, context: ExportContext, output_path: str,
options: Optional[dict] = None) -> ExportResult:
options = options or {}
warnings = []
tokenization = options.get("tokenization", "whitespace")
schema_name = options.get("schema_name")
if not schema_name:
for s in context.schemas:
if s.get("annotation_type") == "span":
schema_name = s.get("name")
break
os.makedirs(output_path, exist_ok=True)
out_file = os.path.join(output_path, "annotations.conllu")
lines = []
total_tokens = 0
total_entities = 0
sent_counter = 0
item_props = context.config.get("item_properties", {})
text_key = item_props.get("text_key", "text")
# Deduplicate by instance
instance_annotations = {}
for ann in context.annotations:
iid = ann.get("instance_id", "")
if iid not in instance_annotations:
instance_annotations[iid] = ann
for instance_id, ann in instance_annotations.items():
item = context.items.get(instance_id, {})
text = item.get(text_key, "")
if not text:
for alt_key in ("text", "sentence", "content"):
if alt_key in item:
text = item[alt_key]
break
if not text:
warnings.append(f"No text found for {instance_id}")
continue
if isinstance(text, list):
text = " ".join(str(t) for t in text)
tokens = tokenize_text(text, method=tokenization)
if not tokens:
continue
# Get spans
spans = []
for span_schema, span_list in ann.get("spans", {}).items():
if schema_name and span_schema != schema_name:
continue
for sp in span_list:
spans.append({
"start": sp.get("start", 0),
"end": sp.get("end", 0),
"label": sp.get("name") or sp.get("label", "ENTITY"),
})
bio_tags = char_spans_to_bio_tags(tokens, spans)
total_tokens += len(tokens)
total_entities += sum(1 for t in bio_tags if t.startswith("B-"))
sentences = group_sentences(tokens, text)
for sentence_indices in sentences:
sent_counter += 1
sent_tokens = [tokens[i] for i in sentence_indices]
sent_tags = [bio_tags[i] for i in sentence_indices]
# Reconstruct sentence text
if sent_tokens:
s_start = sent_tokens[0]["start"]
s_end = sent_tokens[-1]["end"]
sent_text = text[s_start:s_end]
else:
sent_text = ""
lines.append(f"# sent_id = {instance_id}-s{sent_counter}")
lines.append(f"# text = {sent_text}")
for tok_num, (tok, ner_tag) in enumerate(
zip(sent_tokens, sent_tags), start=1
):
# Build MISC field
misc_parts = []
# SpaceAfter=No if no space before next token
if tok_num < len(sent_tokens):
next_tok = sent_tokens[tok_num] # 0-indexed next
if tok["end"] == next_tok["start"]:
misc_parts.append("SpaceAfter=No")
# NER tag
if ner_tag != "O":
misc_parts.append(f"NER={ner_tag}")
misc = "|".join(misc_parts) if misc_parts else "_"
# 10-column CoNLL-U format
# ID FORM LEMMA UPOS XPOS FEATS HEAD DEPREL DEPS MISC
cols = [
str(tok_num), # ID
tok["token"], # FORM
"_", # LEMMA
"_", # UPOS
"_", # XPOS
"_", # FEATS
"_", # HEAD
"_", # DEPREL
"_", # DEPS
misc, # MISC
]
lines.append("\t".join(cols))
lines.append("") # Blank line between sentences
with open(out_file, "w") as f:
f.write("\n".join(lines))
return ExportResult(
success=True,
format_name=self.format_name,
files_written=[out_file],
warnings=warnings,
stats={
"num_documents": len(instance_annotations),
"num_sentences": sent_counter,
"num_tokens": total_tokens,
"num_entities": total_entities,
},
)