You can optimize the TypedEntityRelationshipExtractor from nano_graphrag using the DSPy framework to improve extraction performance. The process involves evaluating a baseline model, then using optimizers like BootstrapFewShotWithRandomSearch or MIPROv2 to refine prompt instructions and few-shot examples based on specific metrics.
Workflow
- Load Data: Prepare training, validation, and development datasets (typically as pickled objects).
- Baseline Evaluation: Use
dspy.evaluate.Evaluate with metrics like entity_recall_metric and relationships_similarity_metric to establish a baseline score for the TypedEntityRelationshipExtractor. - Optimization:
- Use
BootstrapFewShotWithRandomSearch for simple bootstrapping. - Use
MIPROv2 for advanced optimization that generates candidate instructions and few-shot examples using a larger 'prompt model' to guide a smaller 'task model'.
- Save Model: Once optimized, save the compiled model using the
.save() method.
from nano_graphrag.entity_extraction.module import TypedEntityRelationshipExtractor
from nano_graphrag.entity_extraction.metric import relationships_similarity_metric, entity_recall_metric
import dspy
# Initialize the extractor
model = TypedEntityRelationshipExtractor()
# Example: Using BootstrapFewShotWithRandomSearch
optimizer = dspy.teleprompt.BootstrapFewShotWithRandomSearch(
metric=relationships_similarity_metric,
num_threads=os.cpu_count(),
num_candidate_programs=10,
max_labeled_demos=5,
max_bootstrapped_demos=2,
)
rs_model = optimizer.compile(model, trainset=trainset, valset=valset)
# Save the optimized model
rs_model.save("path_to_save_model.json")