""" TRPG CoNLL 转 Dataset 工具 - 自动检测 word-level / char-level - 生成 {"text": str, "char_labels": List[str]} - 支持多文档、跨行实体 """ import os import re import json import argparse from pathlib import Path from typing import List, Dict, Any, Tuple from datasets import Dataset def word_to_char_labels(text: str, word_labels: List[Tuple[str, str]]) -> List[str]: """ 将 word-level 标注转为 char-level labels Args: text: 原始文本 (e.g., "风雨 2024-06-08") word_labels: [("风雨", "B-speaker"), ("2024-06-08", "B-timestamp"), ...] Returns: char_labels: ["B-speaker", "I-speaker", "O", "B-timestamp", ...] """ char_labels = ["O"] * len(text) pos = 0 for token, label in word_labels: if pos >= len(text): break # 在文本中定位 token（处理空格/换行） while pos < len(text) and text[pos] != token[0]: pos += 1 if pos >= len(text): break # 匹配 token if text[pos:pos+len(token)] == token: # 标注 B/I for i, char in enumerate(token): idx = pos + i if idx < len(char_labels): if i == 0 and label.startswith("B-"): char_labels[idx] = label elif label.startswith("B-"): char_labels[idx] = "I" + label[1:] else: char_labels[idx] = label pos += len(token) else: pos += 1 return char_labels def parse_conll_to_samples(filepath: str) -> List[Dict[str, Any]]: """ 解析 .conll → [{"text": "...", "char_labels": [...]}, ...] 自动处理： - -DOCSTART- 文档边界 - 空行句子边界 - word-level → char-level 转换 """ samples = [] current_lines = [] # 存储原始行用于检测粒度 with open(filepath, 'r', encoding='utf-8') as f: for line in f: current_lines.append(line.rstrip('\n')) # 检测是否 word-level is_word_level = False for line in current_lines: if line.strip() and not line.startswith("-DOCSTART-"): parts = line.split() if len(parts) >= 4: token = parts[0] # 如果 token 长度 >1 且非标点 → 可能是 word-level if len(token) > 1 and not re.match(r'^[^\w\s\u4e00-\u9fff]+$', token): is_word_level = True break if is_word_level: print(f"Detected word-level CoNLL, converting to char-level...") return _parse_word_conll(filepath) else: print(f"Detected char-level CoNLL, parsing directly...") return _parse_char_conll(filepath) def _parse_word_conll(filepath: str) -> List[Dict[str, Any]]: """解析 word-level .conll（如您提供的原始格式）""" samples = [] current_text_parts = [] current_word_labels = [] with open(filepath, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line or line.startswith("-DOCSTART-"): if current_text_parts: # 合并文本 text = "".join(current_text_parts) # 生成 char-level labels char_labels = word_to_char_labels(text, current_word_labels) samples.append({ "text": text, "char_labels": char_labels }) current_text_parts = [] current_word_labels = [] continue parts = line.split() if len(parts) < 4: continue token, label = parts[0], parts[3] current_text_parts.append(token) current_word_labels.append((token, label)) # 处理末尾 if current_text_parts: text = "".join(current_text_parts) char_labels = word_to_char_labels(text, current_word_labels) samples.append({ "text": text, "char_labels": char_labels }) return samples def _parse_char_conll(filepath: str) -> List[Dict[str, Any]]: """解析 char-level .conll""" samples = [] current_text = [] current_labels = [] with open(filepath, 'r', encoding='utf-8') as f: for line in f: line = line.rstrip('\n') if line.startswith("-DOCSTART-"): if current_text: samples.append({ "text": "".join(current_text), "char_labels": current_labels.copy() }) current_text, current_labels = [], [] continue if not line: if current_text: samples.append({ "text": "".join(current_text), "char_labels": current_labels.copy() }) current_text, current_labels = [], [] continue parts = line.split() if len(parts) < 4: continue char = parts[0].replace("\\n", "\n") label = parts[3] current_text.append(char) current_labels.append(label) # 末尾处理 if current_text: samples.append({ "text": "".join(current_text), "char_labels": current_labels.copy() }) return samples def save_dataset(samples: List[Dict[str, Any]], output_path: str, format: str = "jsonl"): """保存数据集""" Path(output_path).parent.mkdir(parents=True, exist_ok=True) if format == "jsonl": with open(output_path, 'w', encoding='utf-8') as f: for sample in samples: f.write(json.dumps(sample, ensure_ascii=False) + '\n') print(f"Saved {len(samples)} samples to {output_path} (JSONL)") elif format == "dataset": dataset = Dataset.from_list(samples) dataset.save_to_disk(output_path) print(f"Saved {len(samples)} samples to {output_path} (Hugging Face Dataset)") elif format == "both": jsonl_path = output_path + ".jsonl" with open(jsonl_path, 'w', encoding='utf-8') as f: for sample in samples: f.write(json.dumps(sample, ensure_ascii=False) + '\n') print(f"Saved JSONL to {jsonl_path}") dataset_path = output_path + "_dataset" dataset = Dataset.from_list(samples) dataset.save_to_disk(dataset_path) print(f"Saved Dataset to {dataset_path}") def validate_samples(samples: List[Dict[str, Any]]) -> bool: """验证样本一致性""" for i, sample in enumerate(samples): if len(sample["text"]) != len(sample["char_labels"]): print(f"Sample {i}: text len={len(sample['text'])}, labels len={len(sample['char_labels'])}") return False print(f"All {len(samples)} samples validated: text & labels length match") return True def main(): parser = argparse.ArgumentParser(description="Convert CoNLL to TRPG Dataset") parser.add_argument("input", type=str, help="Input .conll file or directory") parser.add_argument("--output", type=str, default="./dataset/trpg", help="Output path (without extension)") parser.add_argument("--format", choices=["jsonl", "dataset", "both"], default="jsonl", help="Output format") parser.add_argument("--validate", action="store_true", help="Validate samples after conversion") args = parser.parse_args() filepaths = [] if os.path.isdir(args.input): filepaths = sorted(Path(args.input).glob("*.conll")) elif args.input.endswith(".conll"): filepaths = [Path(args.input)] else: raise ValueError("Input must be .conll file or directory") if not filepaths: raise FileNotFoundError(f"No .conll files found in {args.input}") print(f"Processing {len(filepaths)} files: {[f.name for f in filepaths]}") all_samples = [] for fp in filepaths: print(f"\nProcessing {fp.name}...") samples = parse_conll_to_samples(str(fp)) print(f" → {len(samples)} samples") all_samples.extend(samples) print(f"\nTotal: {len(all_samples)} samples") if args.validate: if not validate_samples(all_samples): exit(1) save_dataset(all_samples, args.output, args.format) label_counts = {} for sample in all_samples: for label in sample["char_labels"]: label_counts[label] = label_counts.get(label, 0) + 1 print("\nLabel distribution:") for label in sorted(label_counts.keys()): print(f" {label}: {label_counts[label]}") if __name__ == "__main__": main()