feat: enhance Hugging Face integration with robust BYOM support and improved triplet/relation extraction

tests/test_hf_deep_verify.py

@@ -0,0 +1,167 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import sys
+import os
+
+# Ensure project root is in path
+sys.path.append(os.getcwd())
+
+from semantica.semantic_extract.ner_extractor import NERExtractor, Entity
+from semantica.semantic_extract.relation_extractor import RelationExtractor
+from semantica.semantic_extract.triplet_extractor import TripletExtractor
+from semantica.utils.exceptions import ProcessingError
+
+class TestHuggingFaceDeepIntegration(unittest.TestCase):
+    """
+    Comprehensive test suite for Hugging Face models integration
+    in NER, Relation, and Triplet extraction modules.
+    """
+
+    @patch('semantica.semantic_extract.methods.HuggingFaceModelLoader')
+    def test_ner_extraction_flow(self, MockLoaderClass):
+        """Test NER extraction with detailed IOB parsing and aggregation."""
+        mock_loader = MockLoaderClass.return_value
+        
+        # Simulate IOB output (Raw token classification)
+        mock_loader.extract_entities.return_value = [
+            {"entity": "B-PER", "score": 0.99, "index": 1, "word": "John", "start": 0, "end": 4, "label": "B-PER"},
+            {"entity": "I-PER", "score": 0.98, "index": 2, "word": "Doe", "start": 5, "end": 8, "label": "I-PER"},
+            {"entity": "O", "score": 0.99, "index": 3, "word": "lives", "start": 9, "end": 14, "label": "O"},
+            {"entity": "B-LOC", "score": 0.95, "index": 4, "word": "New", "start": 18, "end": 21, "label": "B-LOC"},
+            {"entity": "I-LOC", "score": 0.96, "index": 5, "word": "York", "start": 22, "end": 26, "label": "I-LOC"},
+        ]
+        
+        extractor = NERExtractor(method="huggingface", huggingface_model="dslim/bert-base-NER")
+        entities = extractor.extract_entities("John Doe lives in New York")
+        
+        # Verify aggregation worked (John Doe should be one entity)
+        # Note: The logic in extract_entities_huggingface handles manual aggregation
+        # if "entity_group" is missing and labels start with B-/I-
+        
+        # Let's debug what we expect. 
+        # "John" (B-PER) -> current_entity="John"
+        # "Doe" (I-PER) -> match! -> current_entity="John Doe"
+        # "lives" (O) -> append John Doe, current=None
+        # "New" (B-LOC) -> current="New"
+        # "York" (I-LOC) -> match! -> current="New York"
+        # End -> append New York
+        
+        self.assertEqual(len(entities), 2)
+        
+        person = next((e for e in entities if e.label == "PER"), None)
+        self.assertIsNotNone(person)
+        self.assertEqual(person.text, "John Doe")
+        
+        loc = next((e for e in entities if e.label == "LOC"), None)
+        self.assertIsNotNone(loc)
+        self.assertEqual(loc.text, "New York")
+
+    @patch('semantica.semantic_extract.methods.HuggingFaceModelLoader')
+    def test_ner_aggregation_strategy_simple(self, MockLoaderClass):
+        """Test NER extraction when the pipeline handles aggregation (strategy='simple')."""
+        mock_loader = MockLoaderClass.return_value
+        
+        # Simulate Aggregated output
+        mock_loader.extract_entities.return_value = [
+            {"entity_group": "PER", "score": 0.99, "word": "John Doe", "start": 0, "end": 8},
+            {"entity_group": "LOC", "score": 0.95, "word": "New York", "start": 18, "end": 26},
+        ]
+        
+        extractor = NERExtractor(
+            method="huggingface", 
+            huggingface_model="dslim/bert-base-NER",
+            aggregation_strategy="simple" # Explicitly requesting simple
+        )
+        entities = extractor.extract_entities("John Doe lives in New York")
+        
+        self.assertEqual(len(entities), 2)
+        self.assertEqual(entities[0].text, "John Doe")
+        self.assertEqual(entities[0].label, "PER")
+        
+    @patch('semantica.semantic_extract.methods.HuggingFaceModelLoader')
+    def test_relation_extraction_flow(self, MockLoaderClass):
+        """Test Relation extraction with Hugging Face model."""
+        mock_loader = MockLoaderClass.return_value
+        
+        # Mock extract_relations output
+        mock_loader.extract_relations.return_value = [{
+            "subject": Entity(text="Apple", label="ORG", start_char=0, end_char=5),
+            "object": Entity(text="Steve Jobs", label="PERSON", start_char=21, end_char=31),
+            "relation": "founded_by",
+            "score": 0.9
+        }]
+        
+        # We need to provide entities for relation extraction usually
+        entities = [
+            Entity(text="Apple", label="ORG", start_char=0, end_char=5),
+            Entity(text="Steve Jobs", label="PERSON", start_char=21, end_char=31)
+        ]
+        
+        extractor = RelationExtractor(method="huggingface", huggingface_model="facebook/bart-large-mnli")
+        relations = extractor.extract_relations("Apple was founded by Steve Jobs", entities=entities)
+        
+        # Check if relation is found
+        self.assertEqual(len(relations), 1)
+        self.assertEqual(relations[0].predicate, "founded_by")
+        self.assertEqual(relations[0].subject.text, "Apple")
+        self.assertEqual(relations[0].object.text, "Steve Jobs")
+
+    @patch('semantica.semantic_extract.methods.HuggingFaceModelLoader')
+    def test_triplet_extraction_rebel(self, MockLoaderClass):
+        """Test Triplet extraction using REBEL parsing logic."""
+        mock_loader = MockLoaderClass.return_value
+        
+        # Mock extract_triplets output
+        # The extract_triplets method in Loader returns [{"triplet": decoded_text}]
+        # But wait, methods.py extract_triplets_huggingface handles parsing?
+        # No, let's check methods.py again.
+        
+        # Actually, methods.py for triplets calls loader.extract_triplets and then parses the result?
+        # Or does loader.extract_triplets return the raw generation?
+        # Let's check the code I read earlier.
+        # loader.extract_triplets returns [{"triplet": decoded}]
+        
+        # But methods.py `extract_triplets_huggingface` logic needs to be verified.
+        # I didn't read extract_triplets_huggingface in methods.py yet (I read entities).
+        # Assuming standard behavior, let's return what loader returns.
+        
+        mock_loader.extract_triplets.return_value = [{"triplet": "<triplet> Apple <subj> founded by <obj> Steve Jobs"}]
+        
+        # Wait, if methods.py expects raw text and parses it, then I need to know IF methods.py does the parsing or if it expects pre-parsed.
+        # Usually, if it's REBEL, the parsing happens after generation.
+        # Let's assume methods.py parses the REBEL format.
+        
+        extractor = TripletExtractor(method="huggingface", huggingface_model="Babelscape/rebel-large")
+        
+        # If the extractor relies on methods.py to parse, and methods.py relies on REBEL format:
+        triplets = extractor.extract_triplets("Apple was founded by Steve Jobs")
+        
+        # Note: If this fails, it might be because I need to check how extract_triplets_huggingface is implemented.
+        # But let's try.
+        if not triplets:
+             # Fallback: maybe methods.py expects the model to return parsed triplets?
+             pass
+             
+        self.assertTrue(len(triplets) > 0)
+        self.assertEqual(triplets[0].subject, "Apple")
+        self.assertEqual(triplets[0].object, "Steve Jobs")
+        self.assertEqual(triplets[0].predicate, "founded by")
+
+    @patch('semantica.semantic_extract.methods.HuggingFaceModelLoader')
+    def test_byom_override(self, MockLoaderClass):
+        """Verify Bring Your Own Model (runtime override) works for all extractors."""
+        mock_loader = MockLoaderClass.return_value
+        mock_loader.extract_entities.return_value = []
+        
+        # NER
+        ner = NERExtractor(method="huggingface", huggingface_model="default-ner")
+        ner.extract_entities("test", huggingface_model="runtime-ner")
+        
+        # Check if load_ner_model was called with runtime model
+        # mock_loader.load_ner_model.assert_called_with("runtime-ner", ...)
+        # args[0] should be "runtime-ner"
+        call_args = mock_loader.load_ner_model.call_args
+        self.assertEqual(call_args[0][0], "runtime-ner")
+            
+if __name__ == "__main__":
+    unittest.main()