mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
KaifAhmad1
b3a239ccb1
- Implemented 'Bring Your Own Model' (BYOM) support for NER, Relation, and Triplet extraction - Added NER aggregation strategies (simple, max, average) - Implemented Relation Extraction via Sequence Classification with entity markers - Enhanced Triplet Extraction with REBEL post-processing and lazy loading - Updated all extractors to prioritize runtime options over config defaults - Added extensive tests and examples (huggingface_demo.py) - Updated documentation and CHANGELOG
169 lines
7.7 KiB
Python
169 lines
7.7 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
import sys
|
|
import os
|
|
|
|
# Add project root to path
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
|
|
# Mock dependencies to avoid import hangs and external calls
|
|
sys.modules['spacy'] = MagicMock()
|
|
sys.modules['semantica.semantic_extract.methods'] = MagicMock()
|
|
sys.modules['semantica.utils.logging'] = MagicMock()
|
|
sys.modules['semantica.utils.progress_tracker'] = MagicMock()
|
|
sys.modules['semantica.semantic_extract.providers'] = MagicMock()
|
|
|
|
# Mock get_logger and get_progress_tracker
|
|
mock_logger = MagicMock()
|
|
sys.modules['semantica.utils.logging'].get_logger.return_value = mock_logger
|
|
|
|
mock_tracker = MagicMock()
|
|
sys.modules['semantica.utils.progress_tracker'].get_progress_tracker.return_value = mock_tracker
|
|
|
|
# Mock the methods module functions specifically
|
|
mock_methods = sys.modules['semantica.semantic_extract.methods']
|
|
mock_methods.get_entity_method = MagicMock()
|
|
mock_methods.get_relation_method = MagicMock()
|
|
mock_methods.get_triplet_method = MagicMock()
|
|
|
|
# Mock specific extraction functions
|
|
mock_extract_entities_hf = MagicMock()
|
|
mock_extract_relations_hf = MagicMock()
|
|
mock_extract_triplets_hf = MagicMock()
|
|
|
|
# Setup the registry mocks to return our mock functions
|
|
mock_methods.get_entity_method.return_value = mock_extract_entities_hf
|
|
mock_methods.get_relation_method.return_value = mock_extract_relations_hf
|
|
mock_methods.get_triplet_method.return_value = mock_extract_triplets_hf
|
|
|
|
# Now import the classes under test
|
|
# We need to patch where they import 'methods' locally if they do
|
|
with patch.dict(sys.modules):
|
|
from semantica.semantic_extract.ner_extractor import NERExtractor
|
|
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
|
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
|
from semantica.semantic_extract.ner_extractor import Entity
|
|
from semantica.semantic_extract.relation_extractor import Relation
|
|
|
|
class TestExtractorsDispatch(unittest.TestCase):
|
|
def setUp(self):
|
|
self.mock_extract_entities_hf = mock_extract_entities_hf
|
|
self.mock_extract_relations_hf = mock_extract_relations_hf
|
|
self.mock_extract_triplets_hf = mock_extract_triplets_hf
|
|
|
|
self.mock_extract_entities_hf.reset_mock()
|
|
self.mock_extract_relations_hf.reset_mock()
|
|
self.mock_extract_triplets_hf.reset_mock()
|
|
|
|
# Configure mocks to return something iterable/valid
|
|
self.mock_extract_entities_hf.return_value = [MagicMock(spec=Entity, confidence=0.9, text="Test Entity")]
|
|
self.mock_extract_relations_hf.return_value = [MagicMock(spec=Relation, confidence=0.9)]
|
|
self.mock_extract_triplets_hf.return_value = [MagicMock(confidence=0.9)]
|
|
|
|
def test_ner_extractor_huggingface_dispatch(self):
|
|
print("\nTesting NERExtractor dispatch to HuggingFace...")
|
|
# Initialize with HuggingFace method
|
|
extractor = NERExtractor(method="huggingface")
|
|
|
|
# Call extract_entities
|
|
text = "Steve Jobs founded Apple."
|
|
# Use a specific model via kwargs
|
|
extractor.extract_entities(text, model="my-custom-ner-model")
|
|
|
|
# Verify get_entity_method was called with "huggingface"
|
|
mock_methods.get_entity_method.assert_called_with("huggingface")
|
|
|
|
# Verify the extraction function was called with correct model
|
|
# We need to check the call args to see if 'model' was passed correctly
|
|
# The logic we implemented: method_options["model"] = all_options.get("huggingface_model") or all_options.get("model") or self.huggingface_model
|
|
|
|
call_args = self.mock_extract_entities_hf.call_args
|
|
self.assertIsNotNone(call_args, "extract_entities_huggingface should have been called")
|
|
|
|
_, kwargs = call_args
|
|
self.assertEqual(kwargs.get("model"), "my-custom-ner-model", "Should use model passed in kwargs")
|
|
|
|
print("NERExtractor dispatch verified.")
|
|
|
|
def test_relation_extractor_huggingface_dispatch(self):
|
|
print("\nTesting RelationExtractor dispatch to HuggingFace...")
|
|
extractor = RelationExtractor(method="huggingface")
|
|
|
|
text = "Steve Jobs founded Apple."
|
|
entities = [MagicMock(spec=Entity)]
|
|
|
|
# Call extract_relations with explicit model
|
|
extractor.extract_relations(text, entities, model="my-relation-model")
|
|
|
|
# Verify dispatch
|
|
mock_methods.get_relation_method.assert_called_with("huggingface")
|
|
|
|
call_args = self.mock_extract_relations_hf.call_args
|
|
self.assertIsNotNone(call_args, "extract_relations_huggingface should have been called")
|
|
|
|
_, kwargs = call_args
|
|
self.assertEqual(kwargs.get("model"), "my-relation-model", "Should use model passed in kwargs")
|
|
|
|
print("RelationExtractor dispatch verified.")
|
|
|
|
def test_triplet_extractor_huggingface_dispatch(self):
|
|
print("\nTesting TripletExtractor dispatch to HuggingFace...")
|
|
extractor = TripletExtractor(method="huggingface")
|
|
|
|
text = "Steve Jobs founded Apple."
|
|
|
|
# Call extract_triplets with explicit model
|
|
extractor.extract_triplets(text, model="my-triplet-model")
|
|
|
|
# Verify dispatch
|
|
mock_methods.get_triplet_method.assert_called_with("huggingface")
|
|
|
|
call_args = self.mock_extract_triplets_hf.call_args
|
|
self.assertIsNotNone(call_args, "extract_triplets_huggingface should have been called")
|
|
|
|
_, kwargs = call_args
|
|
self.assertEqual(kwargs.get("model"), "my-triplet-model", "Should use model passed in kwargs")
|
|
|
|
print("TripletExtractor dispatch verified.")
|
|
|
|
def test_ner_extractor_huggingface_fallback(self):
|
|
print("\nTesting NERExtractor fallback logic...")
|
|
# Init with huggingface_model in config
|
|
extractor = NERExtractor(method="huggingface", huggingface_model="config-model")
|
|
|
|
extractor.extract_entities("text")
|
|
|
|
_, kwargs = self.mock_extract_entities_hf.call_args
|
|
self.assertEqual(kwargs.get("model"), "config-model", "Should prioritize huggingface_model from config")
|
|
|
|
# Now override with kwargs model
|
|
extractor.extract_entities("text", model="kwargs-model")
|
|
_, kwargs = self.mock_extract_entities_hf.call_args
|
|
self.assertEqual(kwargs.get("model"), "kwargs-model", "Should allow overriding config huggingface_model via model kwarg")
|
|
|
|
# Let's test passing 'huggingface_model' in kwargs
|
|
extractor.extract_entities("text", huggingface_model="override-model")
|
|
_, kwargs = self.mock_extract_entities_hf.call_args
|
|
self.assertEqual(kwargs.get("model"), "override-model", "Should allow overriding huggingface_model via kwargs")
|
|
|
|
def test_triplet_extractor_lazy_loading(self):
|
|
print("\nTesting TripletExtractor lazy loading for HuggingFace...")
|
|
# Initialize with HuggingFace method
|
|
extractor = TripletExtractor(method="huggingface")
|
|
|
|
# Check initial state
|
|
self.assertIsNone(extractor._ner_extractor)
|
|
self.assertIsNone(extractor._relation_extractor)
|
|
|
|
# Run extraction
|
|
extractor.extract_triplets("Steve Jobs founded Apple.")
|
|
|
|
# Check state AFTER extraction - should STILL be None because huggingface (REBEL) doesn't need them
|
|
self.assertIsNone(extractor._ner_extractor, "NERExtractor should not be initialized for HuggingFace method")
|
|
self.assertIsNone(extractor._relation_extractor, "RelationExtractor should not be initialized for HuggingFace method")
|
|
|
|
print("TripletExtractor lazy loading verified.")
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|