Files

33 lines
986 B
Python

import pytest
from semantica.context.context_graph import ContextGraph
from semantica.context.entity_linker import EntityLinker
@pytest.mark.benchmark(group="entity_linkiing")
@pytest.mark.parametrize("num_entities_in_graph", [100, 1000])
def test_entity_linking_complexity(benchmark, num_entities_in_graph):
"""
Benchmarks finding links for extracted entities
against the existing graph.
"""
graph = ContextGraph()
nodes = [
{"id": f"e_{i}", "type": "Entity", "properties": {"content": f"Entity {i}"}}
for i in range(num_entities_in_graph)
]
graph.add_nodes(nodes)
graph_dict = graph.to_dict()
linker = EntityLinker(knowledge_graph=graph_dict, similarity_threshold=0.7)
# Simulate extraction
extracted_entities = [{"text": f"Entity {i}", "type": "Entity"} for i in range(5)]
def run():
return linker.link("dummy text", entities=extracted_entities)
benchmark.pedantic(run, iterations=1, rounds=5)