265 lines
8.3 KiB
Python
265 lines
8.3 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Phase 5 Entity Resolver tests."""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
sys.path.insert(0, str(Path(__file__).parent / "ontology_platform"))
|
||
|
|
|
||
|
|
from ont_platform.core.graph.entity_resolver import EntityResolver, EntityCluster
|
||
|
|
|
||
|
|
|
||
|
|
async def test_normalize_label():
|
||
|
|
"""Test label normalization."""
|
||
|
|
print("\n[TEST 1] Label Normalization")
|
||
|
|
|
||
|
|
resolver = EntityResolver()
|
||
|
|
|
||
|
|
test_cases = [
|
||
|
|
("iPhone Pro Max", "iphone pro max"),
|
||
|
|
("The Apple Inc.", "apple inc"),
|
||
|
|
("Test-Entity", "test entity"),
|
||
|
|
("UPPERCASE LABEL", "uppercase label"),
|
||
|
|
("Label with spaces", "label with spaces"),
|
||
|
|
]
|
||
|
|
|
||
|
|
for input_label, expected in test_cases:
|
||
|
|
result = resolver._normalize_label(input_label)
|
||
|
|
status = "[OK]" if result == expected else "[FAIL]"
|
||
|
|
print(f" {status} '{input_label}' -> '{result}' (expected: '{expected}')")
|
||
|
|
assert result == expected, f"Expected '{expected}', got '{result}'"
|
||
|
|
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_jaro_winkler_similarity():
|
||
|
|
"""Test Jaro-Winkler text similarity."""
|
||
|
|
print("\n[TEST 2] Jaro-Winkler Similarity")
|
||
|
|
|
||
|
|
resolver = EntityResolver()
|
||
|
|
|
||
|
|
test_cases = [
|
||
|
|
("iphone", "iphone", 1.0), # Exact match
|
||
|
|
("iphone", "iPhone", None), # Will be normalized before comparison
|
||
|
|
("apple", "aplicant", None), # Similar but not identical
|
||
|
|
("test", "best", None), # Partial match
|
||
|
|
]
|
||
|
|
|
||
|
|
for s1, s2, expected_range in test_cases:
|
||
|
|
sim = resolver._jaro_winkler_similarity(s1, s2)
|
||
|
|
print(f" Similarity('{s1}', '{s2}') = {sim:.3f}")
|
||
|
|
|
||
|
|
if expected_range == 1.0:
|
||
|
|
assert sim == 1.0, f"Expected 1.0, got {sim}"
|
||
|
|
elif expected_range == 0.0:
|
||
|
|
assert sim == 0.0, f"Expected 0.0, got {sim}"
|
||
|
|
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_text_similarity():
|
||
|
|
"""Test combined text similarity."""
|
||
|
|
print("\n[TEST 3] Text Similarity (Jaro-Winkler + Token Overlap)")
|
||
|
|
|
||
|
|
resolver = EntityResolver()
|
||
|
|
|
||
|
|
test_cases = [
|
||
|
|
("machine learning", "machine learning", 1.0),
|
||
|
|
("machine learning", "learning machine", 0.6), # Same tokens, different order
|
||
|
|
("apple", "apple inc", 0.5), # Partial match
|
||
|
|
("test", "best", 0.4), # Phonetically similar
|
||
|
|
]
|
||
|
|
|
||
|
|
for label1, label2, min_expected in test_cases:
|
||
|
|
sim = resolver._compute_text_similarity(label1, label2)
|
||
|
|
status = "[OK]" if sim >= min_expected else "[FAIL]"
|
||
|
|
print(f" {status} TextSim('{label1}', '{label2}') = {sim:.3f} (>= {min_expected})")
|
||
|
|
assert sim >= min_expected, f"Expected >= {min_expected}, got {sim}"
|
||
|
|
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_embedder_initialization():
|
||
|
|
"""Test embedding model initialization."""
|
||
|
|
print("\n[TEST 4] Embedder Initialization")
|
||
|
|
|
||
|
|
resolver = EntityResolver(model_name="all-MiniLM-L6-v2")
|
||
|
|
|
||
|
|
success = await resolver.initialize_embedder()
|
||
|
|
assert success, "Failed to initialize embedder"
|
||
|
|
|
||
|
|
assert resolver.embedder is not None, "Embedder not loaded"
|
||
|
|
print(" [OK] Embedder loaded successfully")
|
||
|
|
|
||
|
|
# Test embedding computation
|
||
|
|
texts = ["machine learning", "artificial intelligence"]
|
||
|
|
embeddings = resolver._embed_batch(texts)
|
||
|
|
|
||
|
|
assert len(embeddings) == 2, f"Expected 2 embeddings, got {len(embeddings)}"
|
||
|
|
assert len(embeddings[0]) == 384, f"Expected 384-dim vectors, got {len(embeddings[0])}-dim"
|
||
|
|
|
||
|
|
print(" [OK] Generated 384-dim embeddings for 2 texts")
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_detect_duplicates():
|
||
|
|
"""Test duplicate detection with vector similarity."""
|
||
|
|
print("\n[TEST 5] Duplicate Detection (Vector + Text)")
|
||
|
|
|
||
|
|
resolver = EntityResolver(
|
||
|
|
vector_threshold=0.85,
|
||
|
|
text_threshold=0.88,
|
||
|
|
)
|
||
|
|
|
||
|
|
success = await resolver.initialize_embedder()
|
||
|
|
assert success, "Failed to initialize embedder"
|
||
|
|
|
||
|
|
# Create test entities with intentional duplicates
|
||
|
|
entities = [
|
||
|
|
{"id": 1, "label": "Apple Inc.", "type": "Company"},
|
||
|
|
{"id": 2, "label": "Apple Inc", "type": "Company"}, # Duplicate (slightly different)
|
||
|
|
{"id": 3, "label": "Microsoft", "type": "Company"},
|
||
|
|
{"id": 4, "label": "Microsoft Corp", "type": "Company"}, # Duplicate
|
||
|
|
{"id": 5, "label": "Google", "type": "Company"},
|
||
|
|
]
|
||
|
|
|
||
|
|
clusters = await resolver.detect_duplicates(entities)
|
||
|
|
|
||
|
|
print(f" Detected {len(clusters)} duplicate clusters")
|
||
|
|
for cluster in clusters:
|
||
|
|
print(
|
||
|
|
f" Cluster: {cluster.canonical_id} ← {cluster.duplicates} "
|
||
|
|
f"(confidence: {cluster.confidence:.3f}, reason: {cluster.reason})"
|
||
|
|
)
|
||
|
|
|
||
|
|
# We expect to find some duplicates
|
||
|
|
assert len(clusters) > 0, "Should detect at least 1 duplicate cluster"
|
||
|
|
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_resolve_cluster():
|
||
|
|
"""Test entity merging."""
|
||
|
|
print("\n[TEST 6] Cluster Resolution (Entity Merging)")
|
||
|
|
|
||
|
|
resolver = EntityResolver()
|
||
|
|
|
||
|
|
# Create test entities
|
||
|
|
entities_map = {
|
||
|
|
1: {
|
||
|
|
"id": 1,
|
||
|
|
"label": "Apple Inc.",
|
||
|
|
"type": "Company",
|
||
|
|
"aliases": ["Apple"],
|
||
|
|
"evidence": [{"text": "Founded in 1976"}],
|
||
|
|
},
|
||
|
|
2: {
|
||
|
|
"id": 2,
|
||
|
|
"label": "Apple",
|
||
|
|
"type": "Company",
|
||
|
|
"aliases": ["AAPL"],
|
||
|
|
"evidence": [{"text": "Technology company"}],
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
cluster = EntityCluster(
|
||
|
|
cluster_id="C_1_2",
|
||
|
|
canonical_id=1,
|
||
|
|
duplicates=[2],
|
||
|
|
confidence=0.92,
|
||
|
|
reason="combined",
|
||
|
|
metadata={},
|
||
|
|
)
|
||
|
|
|
||
|
|
merged = await resolver.resolve_cluster(cluster, entities_map)
|
||
|
|
|
||
|
|
assert merged["id"] == 1, "Canonical ID should be preserved"
|
||
|
|
assert 2 in merged["merged_from"], "Should record merged_from"
|
||
|
|
assert len(merged["aliases"]) >= 3, f"Should consolidate aliases (got {len(merged['aliases'])})"
|
||
|
|
assert len(merged["evidence"]) >= 2, "Should consolidate evidence"
|
||
|
|
|
||
|
|
print(f" [OK] Merged entity with {len(merged['aliases'])} aliases, {len(merged['evidence'])} evidence")
|
||
|
|
print(f" [OK] Aliases: {merged['aliases']}")
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_resolution_report():
|
||
|
|
"""Test resolution report generation."""
|
||
|
|
print("\n[TEST 7] Resolution Report")
|
||
|
|
|
||
|
|
resolver = EntityResolver()
|
||
|
|
|
||
|
|
clusters = [
|
||
|
|
EntityCluster(
|
||
|
|
cluster_id="C_1",
|
||
|
|
canonical_id=1,
|
||
|
|
duplicates=[2, 3],
|
||
|
|
confidence=0.90,
|
||
|
|
reason="combined",
|
||
|
|
metadata={},
|
||
|
|
),
|
||
|
|
EntityCluster(
|
||
|
|
cluster_id="C_2",
|
||
|
|
canonical_id=4,
|
||
|
|
duplicates=[5],
|
||
|
|
confidence=0.85,
|
||
|
|
reason="vector_similarity",
|
||
|
|
metadata={},
|
||
|
|
),
|
||
|
|
]
|
||
|
|
|
||
|
|
report = resolver.get_resolution_report(clusters)
|
||
|
|
|
||
|
|
assert report["total_clusters"] == 2, "Should have 2 clusters"
|
||
|
|
assert report["total_duplicates"] == 3, "Should have 3 total duplicates (2+1)"
|
||
|
|
assert "combined" in report["by_reason"], "Should track reason types"
|
||
|
|
|
||
|
|
print(f" Total clusters: {report['total_clusters']}")
|
||
|
|
print(f" Total duplicates: {report['total_duplicates']}")
|
||
|
|
print(f" Avg confidence: {report['avg_confidence']:.3f}")
|
||
|
|
print(f" By reason: {report['by_reason']}")
|
||
|
|
print(" [PASS]")
|
||
|
|
|
||
|
|
|
||
|
|
async def main():
|
||
|
|
"""Run all tests."""
|
||
|
|
print("=" * 70)
|
||
|
|
print("Phase 5 Entity Resolver Tests")
|
||
|
|
print("=" * 70)
|
||
|
|
|
||
|
|
try:
|
||
|
|
await test_normalize_label()
|
||
|
|
await test_jaro_winkler_similarity()
|
||
|
|
await test_text_similarity()
|
||
|
|
await test_embedder_initialization()
|
||
|
|
await test_detect_duplicates()
|
||
|
|
await test_resolve_cluster()
|
||
|
|
await test_resolution_report()
|
||
|
|
|
||
|
|
print("\n" + "=" * 70)
|
||
|
|
print("All tests passed!")
|
||
|
|
print("=" * 70)
|
||
|
|
print("\nPhase 5.0 Entity Resolver capabilities:")
|
||
|
|
print(" [OK] Label normalization")
|
||
|
|
print(" [OK] Jaro-Winkler text similarity")
|
||
|
|
print(" [OK] Vector embeddings (all-MiniLM-L6-v2)")
|
||
|
|
print(" [OK] Duplicate detection (vector + text)")
|
||
|
|
print(" [OK] Entity merging and consolidation")
|
||
|
|
print(" [OK] Resolution reporting")
|
||
|
|
|
||
|
|
return True
|
||
|
|
except AssertionError as e:
|
||
|
|
print(f"\nTest failed: {e}")
|
||
|
|
return False
|
||
|
|
except Exception as e:
|
||
|
|
print(f"\nUnexpected error: {e}")
|
||
|
|
import traceback
|
||
|
|
traceback.print_exc()
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
success = asyncio.run(main())
|
||
|
|
sys.exit(0 if success else 1)
|