新增修改

final/scripts/export_triplets.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""
+从达梦数据库导出三元组数据到 JSON 文件
+三元组来源：RELATION_INSTANCE 表 (HEAD_ID, RELATION_TYPE, TAIL_ID)
+支持可选清洗：去重、过滤悬空引用、移除空值、规范化关系类型
+用法: python scripts/export_triplets.py
+需设置环境变量 DM_PASSWORD，且 config 中 source=db
+"""
+import json
+import os
+import sys
+
+_script_dir = os.path.dirname(os.path.abspath(__file__))
+_project_root = os.path.dirname(_script_dir)
+sys.path.insert(0, _project_root)
+
+
+def _json_serializer(obj):
+    """处理 Timestamp、datetime、numpy 等不可直接 JSON 序列化的类型"""
+    if hasattr(obj, 'isoformat'):
+        return obj.isoformat()
+    if hasattr(obj, 'item'):
+        return obj.item()
+    if isinstance(obj, float) and (obj != obj or obj == float('inf') or obj == float('-inf')):
+        return None
+    raise TypeError(f'Object of type {type(obj).__name__} is not JSON serializable')
+
+
+def load_config(project_root: str) -> dict:
+    """加载 config"""
+    config_dir = os.path.join(project_root, 'config')
+    json_path = os.path.join(config_dir, 'config.json')
+    yaml_path = os.path.join(config_dir, 'config.yaml')
+    if os.path.exists(json_path):
+        with open(json_path, 'r', encoding='utf-8') as f:
+            return json.load(f)
+    if os.path.exists(yaml_path):
+        try:
+            import yaml
+            with open(yaml_path, 'r', encoding='utf-8') as f:
+                return yaml.safe_load(f)
+        except ImportError:
+            pass
+    return {'source': 'json', 'db': {}, 'triplet': {}}
+
+
+def main():
+    from src.data_loader import load_triplets_from_dm, load_from_dm
+    from src.triplet_cleaner import clean_triplets, transform_triplet_output_format
+
+    config = load_config(_project_root)
+    if config.get('source') != 'db':
+        print("错误：三元组导出仅支持达梦模式，请将 config 中 source 改为 db")
+        sys.exit(1)
+
+    tc = config.get('triplet', {})
+    output_file = tc.get('output_file', 'data/triples.json')
+    report_file = tc.get('report_file', 'report/triplet_cleaning_report.json')
+    output_path = output_file if os.path.isabs(output_file) else os.path.join(_project_root, output_file)
+    report_path = report_file if os.path.isabs(report_file) else os.path.join(_project_root, report_file)
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    os.makedirs(os.path.dirname(report_path), exist_ok=True)
+
+    # 清洗配置（clean: false 则跳过；否则按子项配置）
+    clean_cfg = tc.get('clean', True)
+    do_clean = False
+    if clean_cfg is True:
+        do_clean = {'deduplicate': True, 'filter_orphans': False, 'remove_null_core': True, 'normalize_relation_type': True}
+    elif isinstance(clean_cfg, dict):
+        do_clean = clean_cfg
+
+    deduplicate = do_clean.get('deduplicate', True) if do_clean else False
+    filter_orphans = do_clean.get('filter_orphans', False) if do_clean else False
+    remove_null_core = do_clean.get('remove_null_core', True) if do_clean else False
+    normalize_relation_type = do_clean.get('normalize_relation_type', True) if do_clean else False
+
+    print("正在从达梦读取三元组 (RELATION_INSTANCE)...")
+    triplets = load_triplets_from_dm(config, _project_root)
+    print(f"读取到 {len(triplets)} 条三元组")
+
+    valid_ids = None
+    if filter_orphans:
+        print("正在加载实体表以校验悬空引用...")
+        entities = load_from_dm(config, _project_root)
+        valid_ids = {str(e.get('TARGET_ID', '')).strip() for e in entities if e.get('TARGET_ID')}
+        print(f"有效实体数: {len(valid_ids)}")
+
+    if do_clean and (deduplicate or filter_orphans or remove_null_core or normalize_relation_type):
+        print("正在执行三元组清洗...")
+        triplets, report = clean_triplets(
+            triplets,
+            valid_ids=valid_ids,
+            deduplicate=deduplicate,
+            filter_orphans=filter_orphans,
+            remove_null_core=remove_null_core,
+            normalize_relation_type=normalize_relation_type,
+        )
+        with open(report_path, 'w', encoding='utf-8') as f:
+            json.dump(report, f, ensure_ascii=False, indent=2, default=_json_serializer)
+        print(f"清洗完成：{report['original_count']} -> {report['final_count']}，报告: {report_path}")
+
+    # 转为输出格式：head, head_type, relation, tail, tail_type
+    entities = load_from_dm(config, _project_root)
+    id_to_role = {str(e.get('TARGET_ID', '')).strip(): str(e.get('ROLE_ID', '')).strip() for e in entities if e.get('TARGET_ID')}
+    triplets = transform_triplet_output_format(triplets, id_to_role)
+
+    with open(output_path, 'w', encoding='utf-8') as f:
+        json.dump(triplets, f, ensure_ascii=False, indent=2, default=_json_serializer)
+    print(f"已导出至 {output_path}")
+
+
+if __name__ == "__main__":
+    main()