189 lines
6.5 KiB
Python
189 lines
6.5 KiB
Python
|
#!/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()
|