189 lines
6.5 KiB
Python
Executable File
189 lines
6.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
对比Gold Standard和自动提取结果
|
||
计算Precision, Recall, F1
|
||
"""
|
||
|
||
import json
|
||
import argparse
|
||
from typing import Dict, List, Tuple, Set
|
||
from dataclasses import dataclass
|
||
|
||
|
||
@dataclass
|
||
class ComparisonResult:
|
||
"""对比结果"""
|
||
true_positives: int
|
||
false_positives: int
|
||
false_negatives: int
|
||
precision: float
|
||
recall: float
|
||
f1: float
|
||
|
||
|
||
def normalize_entity(entity: Dict) -> str:
|
||
"""标准化实体用于对比（忽略id，比较type+name+species）"""
|
||
entity_type = (entity.get("type") or "").lower()
|
||
name = (entity.get("name") or entity.get("id") or "").lower().strip()
|
||
species = (entity.get("species") or "").lower().strip()
|
||
|
||
# For genes/proteins, species is important
|
||
if entity_type in ["gene", "protein"]:
|
||
return f"{entity_type}:{name}:{species}"
|
||
else:
|
||
return f"{entity_type}:{name}"
|
||
|
||
|
||
def normalize_relation(relation: Dict, entities: List[Dict]) -> str:
|
||
"""标准化关系用于对比（直接使用ID，不转换名称）"""
|
||
# 支持两种字段名: "type" 或 "relation"
|
||
rel_type = relation.get("type") or relation.get("relation", "")
|
||
rel_type = rel_type.lower()
|
||
source = relation.get("source", "").lower()
|
||
target = relation.get("target", "").lower()
|
||
|
||
# 直接使用ID，不转换名称（避免名称不一致问题）
|
||
return f"{source}--{rel_type}--{target}"
|
||
|
||
|
||
def compare_entities(gold: List[Dict], pred: List[Dict]) -> ComparisonResult:
|
||
"""对比实体"""
|
||
gold_set = {normalize_entity(e) for e in gold}
|
||
pred_set = {normalize_entity(e) for e in pred}
|
||
|
||
tp = len(gold_set & pred_set)
|
||
fp = len(pred_set - gold_set)
|
||
fn = len(gold_set - pred_set)
|
||
|
||
precision = tp / (tp + fp) if (tp + fp) > 0 else 0
|
||
recall = tp / (tp + fn) if (tp + fn) > 0 else 0
|
||
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
|
||
|
||
return ComparisonResult(tp, fp, fn, precision, recall, f1)
|
||
|
||
|
||
def compare_relations(gold: List[Dict], gold_entities: List[Dict],
|
||
pred: List[Dict], pred_entities: List[Dict]) -> ComparisonResult:
|
||
"""对比关系"""
|
||
gold_set = {normalize_relation(r, gold_entities) for r in gold}
|
||
pred_set = {normalize_relation(r, pred_entities) for r in pred}
|
||
|
||
tp = len(gold_set & pred_set)
|
||
fp = len(pred_set - gold_set)
|
||
fn = len(gold_set - pred_set)
|
||
|
||
precision = tp / (tp + fp) if (tp + fp) > 0 else 0
|
||
recall = tp / (tp + fn) if (tp + fn) > 0 else 0
|
||
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
|
||
|
||
return ComparisonResult(tp, fp, fn, precision, recall, f1)
|
||
|
||
|
||
def print_comparison(name: str, result: ComparisonResult, gold: List, pred: List):
|
||
"""打印对比结果"""
|
||
print(f"\n{'='*60}")
|
||
print(f"{name} Comparison")
|
||
print(f"{'='*60}")
|
||
print(f"Gold Standard: {len(gold)} items")
|
||
print(f"Predicted: {len(pred)} items")
|
||
print(f"-"*60)
|
||
print(f"True Positives: {result.true_positives}")
|
||
print(f"False Positives: {result.false_positives}")
|
||
print(f"False Negatives: {result.false_negatives}")
|
||
print(f"-"*60)
|
||
print(f"Precision: {result.precision:.3f}")
|
||
print(f"Recall: {result.recall:.3f}")
|
||
print(f"F1 Score: {result.f1:.3f}")
|
||
|
||
|
||
def print_detailed_comparison(gold: List[Dict], pred: List[Dict], item_type: str = "entity"):
|
||
"""打印详细对比（缺失和多余的项）"""
|
||
|
||
if item_type == "entity":
|
||
gold_set = {normalize_entity(e) for e in gold}
|
||
pred_set = {normalize_entity(e) for e in pred}
|
||
|
||
# Create reverse lookup
|
||
gold_map = {normalize_entity(e): e for e in gold}
|
||
pred_map = {normalize_entity(e): e for e in pred}
|
||
else:
|
||
# For relations, we need entity lists
|
||
return # Skip for now
|
||
|
||
missing = gold_set - pred_set
|
||
extra = pred_set - pred_set
|
||
|
||
if missing:
|
||
print(f"\n⚠️ Missing {item_type}s (False Negatives):")
|
||
for item in sorted(missing)[:10]: # Show first 10
|
||
print(f" - {item}")
|
||
if len(missing) > 10:
|
||
print(f" ... and {len(missing) - 10} more")
|
||
|
||
if extra:
|
||
print(f"\n⚠️ Extra {item_type}s (False Positives):")
|
||
for item in sorted(extra)[:10]:
|
||
print(f" - {item}")
|
||
if len(extra) > 10:
|
||
print(f" ... and {len(extra) - 10} more")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Compare Gold Standard with automatic extraction")
|
||
parser.add_argument("--gold", required=True, help="Gold Standard JSON file")
|
||
parser.add_argument("--pred", required=True, help="Predicted/Auto extraction JSON file")
|
||
parser.add_argument("--detailed", action="store_true", help="Show detailed comparison")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Load files
|
||
print(f"Loading Gold Standard: {args.gold}")
|
||
with open(args.gold, 'r', encoding='utf-8') as f:
|
||
gold_data = json.load(f)
|
||
|
||
print(f"Loading Prediction: {args.pred}")
|
||
with open(args.pred, 'r', encoding='utf-8') as f:
|
||
pred_data = json.load(f)
|
||
|
||
# Extract entities and relations
|
||
gold_entities = gold_data.get("entities", [])
|
||
gold_relations = gold_data.get("relations", [])
|
||
|
||
pred_entities = pred_data.get("entities", [])
|
||
pred_relations = pred_data.get("relations", [])
|
||
|
||
# Compare entities
|
||
entity_result = compare_entities(gold_entities, pred_entities)
|
||
print_comparison("Entity", entity_result, gold_entities, pred_entities)
|
||
|
||
if args.detailed:
|
||
print_detailed_comparison(gold_entities, pred_entities, "entity")
|
||
|
||
# Compare relations
|
||
relation_result = compare_relations(gold_relations, gold_entities,
|
||
pred_relations, pred_entities)
|
||
print_comparison("Relation", relation_result, gold_relations, pred_relations)
|
||
|
||
# Overall summary
|
||
print(f"\n{'='*60}")
|
||
print("Summary")
|
||
print(f"{'='*60}")
|
||
print(f"Entity F1: {entity_result.f1:.3f}")
|
||
print(f"Relation F1: {relation_result.f1:.3f}")
|
||
avg_f1 = (entity_result.f1 + relation_result.f1) / 2
|
||
print(f"Average F1: {avg_f1:.3f}")
|
||
|
||
# Quality assessment
|
||
if avg_f1 >= 0.8:
|
||
print("\n✅ Excellent extraction quality")
|
||
elif avg_f1 >= 0.6:
|
||
print("\n⚠️ Good extraction quality, some improvements needed")
|
||
elif avg_f1 >= 0.4:
|
||
print("\n⚠️ Fair extraction quality, significant improvements needed")
|
||
else:
|
||
print("\n❌ Poor extraction quality, major rework needed")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|