diff --git a/tests/test_hf_deep_verify.py b/tests/test_hf_deep_verify.py new file mode 100644 index 00000000..ab0b0fc5 --- /dev/null +++ b/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": " Apple founded by 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()