mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
Improve semantic_extract performance and add Groq LLM smoke tests
This commit is contained in:
@@ -13,6 +13,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- Added `max_workers` configuration parameter (default: 1) to all extractor `extract()` methods, allowing users to tune concurrency based on available CPU cores or API rate limits.
|
||||
- **Parallel Chunking**: Implemented parallel processing for large document chunking in `_extract_entities_chunked` and `_extract_relations_chunked`, significantly reducing latency for long-form text analysis.
|
||||
- **Thread-Safe Progress Tracking**: Enhanced `ProgressTracker` to handle concurrent updates from multiple threads without race conditions during batch processing.
|
||||
- **Semantic Extract Performance & Regression**:
|
||||
- Added edge-case regression suite covering max worker defaults, LLM prompt entity filtering, and extractor reuse.
|
||||
- Added a runnable real-use-case benchmark script for batch latency across `NERExtractor`, `RelationExtractor`, `TripletExtractor`, `EventDetector`, `SemanticAnalyzer`, and `SemanticNetworkExtractor`.
|
||||
- Added Groq LLM smoke tests that exercise LLM-based entities/relations/triplets when `GROQ_API_KEY` is available via environment configuration.
|
||||
|
||||
### Security
|
||||
- **Credential Sanitization**:
|
||||
@@ -29,12 +33,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- **Dependency Resolution**:
|
||||
- Pinned `opentelemetry-api` and `opentelemetry-sdk` to `1.37.0` to resolve pip conflicts.
|
||||
- Updated `protobuf` and `grpcio` constraints for better stability.
|
||||
- **Entity Filtering Scope**:
|
||||
- Removed entity filtering from non-LLM extraction flows to avoid accuracy regressions.
|
||||
- Applied entity downselection only to LLM relation prompt construction, while matching returned entities against the full original entity list.
|
||||
- **Batch Concurrency Defaults**:
|
||||
- Standardized `max_workers` defaulting across `semantic_extract` and tuned for low-latency: ML-backed methods default to single-worker, while pattern/regex/rules/LLM/huggingface methods use a higher parallelism default capped by CPU.
|
||||
- Raised the global `optimization.max_workers` default to 8 for better throughput on batch workloads.
|
||||
|
||||
### Performance
|
||||
- **Bottleneck Optimization (GitHub Issue #186)**:
|
||||
- **Resolved Bottleneck #1 (Sequential Processing)**: Replaced sequential `for` loops with parallel execution for both document-level batches and intra-document chunks.
|
||||
- **Performance Gains**: Achieved **~1.89x speedup** in real-world extraction scenarios (tested with Groq `llama-3.3-70b-versatile` on standard datasets).
|
||||
- **Initialization Optimization**: Refactored test suite to use class-level `setUpClass` for LLM provider initialization, eliminating redundant API client creation overhead.
|
||||
- **Low-Latency Entity Matching**:
|
||||
- Avoided heavyweight embedding stack imports on common matches by improving fast matching heuristics and short-circuiting before embedding similarity.
|
||||
- Optimized entity matching to prioritize exact/substring/word-boundary matches and only fall back to embedding similarity when needed, reducing CPU overhead in LLM relation/triplet mapping.
|
||||
|
||||
|
||||
## [0.2.1] - 2026-01-12
|
||||
|
||||
@@ -40,6 +40,7 @@ License: MIT
|
||||
"""
|
||||
|
||||
import os
|
||||
import multiprocessing
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Any
|
||||
|
||||
@@ -57,7 +58,7 @@ class Config:
|
||||
self._configs["optimization"] = {
|
||||
"enable_cache": True,
|
||||
"cache_size": 1000,
|
||||
"max_workers": 5,
|
||||
"max_workers": 8,
|
||||
"enable_batching": True,
|
||||
"batch_size": 10,
|
||||
"max_tokens_per_batch": 2000
|
||||
@@ -151,3 +152,43 @@ class Config:
|
||||
|
||||
# Global config instance
|
||||
config = Config()
|
||||
|
||||
|
||||
def resolve_max_workers(
|
||||
explicit: Optional[int] = None,
|
||||
local_config: Optional[Dict[str, Any]] = None,
|
||||
methods: Optional[Any] = None,
|
||||
) -> int:
|
||||
def to_int(val: Any, default: int) -> int:
|
||||
try:
|
||||
return int(val)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
if isinstance(methods, str):
|
||||
normalized_methods = [methods]
|
||||
elif isinstance(methods, (list, tuple, set)):
|
||||
normalized_methods = [m for m in methods if isinstance(m, str)]
|
||||
else:
|
||||
normalized_methods = []
|
||||
|
||||
if explicit is not None:
|
||||
value = to_int(explicit, 1)
|
||||
elif local_config and "max_workers" in local_config:
|
||||
value = to_int(local_config.get("max_workers", 1), 1)
|
||||
else:
|
||||
value = to_int(config.get("max_workers", 5), 5)
|
||||
|
||||
if "ml" in normalized_methods and explicit is None and not (local_config and "max_workers" in local_config):
|
||||
value = 1
|
||||
|
||||
if value < 1:
|
||||
value = 1
|
||||
|
||||
cpu_count = multiprocessing.cpu_count() or 1
|
||||
if value > cpu_count:
|
||||
value = cpu_count
|
||||
if value > 32:
|
||||
value = 32
|
||||
|
||||
return value
|
||||
|
||||
@@ -185,7 +185,12 @@ class EventDetector:
|
||||
message=f"Starting batch detection... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
max_workers = kwargs.get("max_workers", self.config.get("max_workers", 1))
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=[self.config.get("ner_method"), self.config.get("relation_method"), self.config.get("method")],
|
||||
)
|
||||
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
@@ -269,17 +274,26 @@ class EventDetector:
|
||||
# Single item
|
||||
return self.detect_events(text, **kwargs)
|
||||
|
||||
def detect_events(self, text: str, **options) -> List[Event]:
|
||||
def detect_events(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Union[List[Event], List[List[Event]]]:
|
||||
"""
|
||||
Detect events in text content.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Detection options
|
||||
|
||||
Returns:
|
||||
list: List of detected events
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
return self.extract(text, pipeline_id=pipeline_id, **options)
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="EventDetector",
|
||||
|
||||
@@ -287,6 +287,12 @@ def find_best_match_index(text: str, candidates: List[str]) -> Tuple[int, float]
|
||||
best_idx = i
|
||||
|
||||
# 2. Substring Match (Fast)
|
||||
word_pat = None
|
||||
try:
|
||||
word_pat = re.compile(rf"\b{re.escape(text_lower)}\b")
|
||||
except Exception:
|
||||
word_pat = None
|
||||
|
||||
for i, cand in enumerate(candidates_lower):
|
||||
if not cand: continue
|
||||
score = 0.0
|
||||
@@ -294,12 +300,18 @@ def find_best_match_index(text: str, candidates: List[str]) -> Tuple[int, float]
|
||||
# Calculate length ratio
|
||||
ratio = min(len(text_lower), len(cand)) / max(len(text_lower), len(cand))
|
||||
score = 0.9 * ratio + 0.1
|
||||
|
||||
if word_pat and word_pat.search(cand):
|
||||
score = max(score, 0.88)
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_idx = i
|
||||
|
||||
# 3. Text Embeddings (High Accuracy Semantic) - Batch Optimized
|
||||
if best_score >= 0.85:
|
||||
return best_idx, float(best_score)
|
||||
|
||||
embedder = get_text_embedder()
|
||||
embedding_idx = -1
|
||||
embedding_score = 0.0
|
||||
@@ -417,6 +429,86 @@ def match_entity(text: str, entities: List[Entity], threshold: float = 0.8) -> O
|
||||
return None
|
||||
|
||||
|
||||
def filter_entities_for_text(
|
||||
text: str,
|
||||
entities: List[Entity],
|
||||
max_keep: int = 80,
|
||||
) -> List[Entity]:
|
||||
if not text or not entities:
|
||||
return []
|
||||
|
||||
if max_keep < 1:
|
||||
return []
|
||||
|
||||
if len(entities) <= max_keep:
|
||||
return entities
|
||||
|
||||
text_lower = text.lower()
|
||||
stop_tokens = {
|
||||
"inc",
|
||||
"incorporated",
|
||||
"corp",
|
||||
"corporation",
|
||||
"co",
|
||||
"company",
|
||||
"ltd",
|
||||
"llc",
|
||||
"plc",
|
||||
"group",
|
||||
"holdings",
|
||||
"limited",
|
||||
"the",
|
||||
"and",
|
||||
"or",
|
||||
"of",
|
||||
"in",
|
||||
"on",
|
||||
"at",
|
||||
"for",
|
||||
"to",
|
||||
"a",
|
||||
"an",
|
||||
}
|
||||
|
||||
seen = set()
|
||||
matched: List[Entity] = []
|
||||
for entity in entities:
|
||||
ent_text = getattr(entity, "text", "")
|
||||
if not ent_text:
|
||||
continue
|
||||
|
||||
key = ent_text.lower().strip()
|
||||
if not key or key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
|
||||
if key in text_lower:
|
||||
matched.append(entity)
|
||||
continue
|
||||
|
||||
tokens = re.findall(r"[a-z0-9]+", key)
|
||||
keep = False
|
||||
for tok in tokens:
|
||||
if len(tok) < 2:
|
||||
continue
|
||||
if tok in stop_tokens:
|
||||
continue
|
||||
if tok in text_lower:
|
||||
keep = True
|
||||
break
|
||||
if keep:
|
||||
matched.append(entity)
|
||||
|
||||
if matched:
|
||||
if len(matched) > max_keep:
|
||||
matched.sort(key=lambda e: len(getattr(e, "text", "")), reverse=True)
|
||||
return matched[:max_keep]
|
||||
return matched
|
||||
|
||||
entities_sorted = sorted(entities, key=lambda e: len(getattr(e, "text", "")), reverse=True)
|
||||
return entities_sorted[:max_keep]
|
||||
|
||||
|
||||
def calculate_weighted_confidence(
|
||||
item_type: str,
|
||||
original_confidence: float,
|
||||
@@ -907,8 +999,8 @@ def _extract_entities_chunked(
|
||||
|
||||
all_entities = []
|
||||
|
||||
# Process chunks in parallel
|
||||
max_workers = kwargs.get("max_workers", 5)
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(explicit=kwargs.get("max_workers"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_chunk = {}
|
||||
@@ -1513,7 +1605,22 @@ def extract_relations_llm(
|
||||
**kwargs
|
||||
)
|
||||
|
||||
entities_str = ", ".join([f"{e.text} ({e.label})" for e in entities])
|
||||
original_entities = entities
|
||||
max_entities_prompt = kwargs.get("max_entities_prompt", kwargs.get("max_entities", 80))
|
||||
try:
|
||||
max_entities_prompt = int(max_entities_prompt)
|
||||
except Exception:
|
||||
max_entities_prompt = 80
|
||||
|
||||
prompt_entities = original_entities
|
||||
if max_entities_prompt > 0 and len(original_entities) > max_entities_prompt:
|
||||
prompt_entities = filter_entities_for_text(
|
||||
text,
|
||||
original_entities,
|
||||
max_keep=max_entities_prompt,
|
||||
)
|
||||
|
||||
entities_str = ", ".join([f"{e.text} ({e.label})" for e in prompt_entities])
|
||||
|
||||
# Use custom relation types if provided
|
||||
relation_types = kwargs.get("relation_types")
|
||||
@@ -1562,8 +1669,8 @@ Entities found in text: {entities_str}"""
|
||||
relations = []
|
||||
for r_out in result_obj.relations:
|
||||
# Find matching entities using hybrid similarity
|
||||
subject_entity = match_entity(r_out.subject, entities)
|
||||
object_entity = match_entity(r_out.object, entities)
|
||||
subject_entity = match_entity(r_out.subject, original_entities)
|
||||
object_entity = match_entity(r_out.object, original_entities)
|
||||
|
||||
if subject_entity and object_entity:
|
||||
relations.append(Relation(
|
||||
@@ -1689,8 +1796,8 @@ def _extract_relations_chunked(
|
||||
|
||||
all_relations = []
|
||||
|
||||
# Process chunks in parallel
|
||||
max_workers = kwargs.get("max_workers", 5)
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(explicit=kwargs.get("max_workers"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_chunk = {}
|
||||
@@ -2093,8 +2200,8 @@ def _extract_triplets_chunked(
|
||||
|
||||
all_triplets = []
|
||||
|
||||
# Process chunks in parallel
|
||||
max_workers = kwargs.get("max_workers", 5)
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(explicit=kwargs.get("max_workers"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_chunk = {}
|
||||
|
||||
@@ -197,8 +197,12 @@ class NERExtractor:
|
||||
message=f"Starting batch extraction... 0/{total_items}"
|
||||
)
|
||||
|
||||
# Determine max_workers
|
||||
max_workers = kwargs.get("max_workers", self.config.get("max_workers", 1))
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=self.method,
|
||||
)
|
||||
|
||||
# Helper function for single item processing
|
||||
def process_item(idx, item):
|
||||
@@ -297,12 +301,18 @@ class NERExtractor:
|
||||
else:
|
||||
return self.extract_entities(text, **kwargs)
|
||||
|
||||
def extract_entities(self, text: str, **options) -> List[Entity]:
|
||||
def extract_entities(
|
||||
self,
|
||||
text: Union[str, List[Dict[str, Any]], List[str]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Union[List[Entity], List[List[Entity]]]:
|
||||
"""
|
||||
Extract named entities from text.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Extraction options:
|
||||
- entity_types: Filter by entity types (list)
|
||||
- min_confidence: Minimum confidence threshold
|
||||
@@ -311,6 +321,9 @@ class NERExtractor:
|
||||
Returns:
|
||||
list: List of extracted entities
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
return self.extract(text, pipeline_id=pipeline_id, **options)
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="NERExtractor",
|
||||
|
||||
@@ -221,8 +221,12 @@ class RelationExtractor:
|
||||
message=f"Starting batch extraction... 0/{min_len}"
|
||||
)
|
||||
|
||||
# Determine max_workers
|
||||
max_workers = kwargs.get("max_workers", self.config.get("max_workers", 1))
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=self.method,
|
||||
)
|
||||
|
||||
def process_item(i, doc_item, ent_item):
|
||||
try:
|
||||
@@ -325,14 +329,19 @@ class RelationExtractor:
|
||||
return []
|
||||
|
||||
def extract_relations(
|
||||
self, text: str, entities: List[Entity], **options
|
||||
) -> List[Relation]:
|
||||
self,
|
||||
text: Union[str, List[Dict[str, Any]], List[str]],
|
||||
entities: Union[List[Entity], List[List[Entity]]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Union[List[Relation], List[List[Relation]]]:
|
||||
"""
|
||||
Extract relations between entities.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
entities: List of extracted entities
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Extraction options:
|
||||
- method: Override method (if not set in __init__)
|
||||
- min_confidence: Minimum confidence threshold
|
||||
@@ -341,6 +350,17 @@ class RelationExtractor:
|
||||
Returns:
|
||||
list: List of extracted relations
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
if entities is None:
|
||||
entities_batch = [[] for _ in range(len(text))]
|
||||
elif isinstance(entities, list) and (not entities):
|
||||
entities_batch = [[] for _ in range(len(text))]
|
||||
elif isinstance(entities, list) and all(isinstance(e, Entity) for e in entities):
|
||||
entities_batch = [entities for _ in range(len(text))]
|
||||
else:
|
||||
entities_batch = entities
|
||||
return self.extract(text, entities_batch, pipeline_id=pipeline_id, **options)
|
||||
|
||||
from .methods import get_relation_method, match_entity
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
|
||||
@@ -148,8 +148,9 @@ class SemanticAnalyzer:
|
||||
)
|
||||
|
||||
try:
|
||||
results = []
|
||||
results = [None] * len(text)
|
||||
total_items = len(text)
|
||||
processed_count = 0
|
||||
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
@@ -165,40 +166,83 @@ class SemanticAnalyzer:
|
||||
message=f"Starting batch analysis... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
for idx, item in enumerate(text):
|
||||
# Prepare arguments for single item
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
|
||||
# Analyze
|
||||
analysis = self.analyze_semantics(doc_text, **kwargs)
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
)
|
||||
|
||||
# Add provenance metadata
|
||||
analysis["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
analysis["document_id"] = item["id"]
|
||||
|
||||
# Also inject into semantic roles if present
|
||||
if "semantic_roles" in analysis:
|
||||
for role in analysis["semantic_roles"]:
|
||||
# role is a dict here because analyze_semantics converts it
|
||||
if "metadata" not in role:
|
||||
role["metadata"] = {}
|
||||
|
||||
role["metadata"]["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
role["metadata"]["document_id"] = item["id"]
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
analysis = self.analyze_semantics(doc_text, **kwargs)
|
||||
|
||||
results.append(analysis)
|
||||
analysis["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
analysis["document_id"] = item["id"]
|
||||
|
||||
# Update progress
|
||||
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||
remaining = total_items - (idx + 1)
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=idx + 1,
|
||||
total=total_items,
|
||||
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining})"
|
||||
if "semantic_roles" in analysis:
|
||||
for role in analysis["semantic_roles"]:
|
||||
if "metadata" not in role:
|
||||
role["metadata"] = {}
|
||||
|
||||
role["metadata"]["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
role["metadata"]["document_id"] = item["id"]
|
||||
|
||||
return idx, analysis
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to analyze item {idx}: {e}")
|
||||
return idx, {"error": str(e), "batch_index": idx}
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_idx = {
|
||||
executor.submit(process_item, idx, item): idx
|
||||
for idx, item in enumerate(text)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
idx, analysis = future.result()
|
||||
results[idx] = analysis
|
||||
processed_count += 1
|
||||
|
||||
should_update = (
|
||||
processed_count % update_interval == 0
|
||||
or processed_count == total_items
|
||||
or processed_count == 1
|
||||
or total_items <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining})"
|
||||
)
|
||||
else:
|
||||
for idx, item in enumerate(text):
|
||||
_, analysis = process_item(idx, item)
|
||||
results[idx] = analysis
|
||||
processed_count += 1
|
||||
|
||||
should_update = (
|
||||
processed_count % update_interval == 0
|
||||
or processed_count == total_items
|
||||
or processed_count == 1
|
||||
or total_items <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining})"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
|
||||
@@ -149,6 +149,9 @@ class SemanticNetworkExtractor:
|
||||
self.config["ner_method"] = method
|
||||
self.config["relation_method"] = method
|
||||
|
||||
self._ner_extractor = None
|
||||
self._relation_extractor = None
|
||||
|
||||
def extract(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
@@ -199,8 +202,11 @@ class SemanticNetworkExtractor:
|
||||
message=f"Starting batch extraction... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
# Determine max_workers
|
||||
max_workers = kwargs.get("max_workers", self.config.get("max_workers", 1))
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
)
|
||||
|
||||
def process_item(idx, item, doc_entities, doc_relations):
|
||||
try:
|
||||
@@ -318,11 +324,12 @@ class SemanticNetworkExtractor:
|
||||
|
||||
def extract_network(
|
||||
self,
|
||||
text: str,
|
||||
entities: Optional[List[Entity]] = None,
|
||||
relations: Optional[List[Relation]] = None,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> SemanticNetwork:
|
||||
) -> Union[SemanticNetwork, List[SemanticNetwork]]:
|
||||
"""
|
||||
Extract semantic network from text.
|
||||
|
||||
@@ -335,6 +342,23 @@ class SemanticNetworkExtractor:
|
||||
Returns:
|
||||
SemanticNetwork: Extracted semantic network
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
entities_batch = entities
|
||||
if entities is not None and isinstance(entities, list) and (not entities or all(isinstance(e, Entity) for e in entities)):
|
||||
entities_batch = [entities for _ in range(len(text))] if entities else [[] for _ in range(len(text))]
|
||||
|
||||
relations_batch = relations
|
||||
if relations is not None and isinstance(relations, list) and (not relations or all(isinstance(r, Relation) for r in relations)):
|
||||
relations_batch = [relations for _ in range(len(text))] if relations else [[] for _ in range(len(text))]
|
||||
|
||||
return self.extract(
|
||||
text,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
pipeline_id=pipeline_id,
|
||||
**options,
|
||||
)
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="SemanticNetworkExtractor",
|
||||
@@ -354,15 +378,16 @@ class SemanticNetworkExtractor:
|
||||
# Pass method if specified
|
||||
if "ner_method" in self.config:
|
||||
ner_config["method"] = self.config["ner_method"]
|
||||
ner = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
entities = ner.extract_entities(text, **options)
|
||||
if self._ner_extractor is None:
|
||||
self._ner_extractor = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
entities = self._ner_extractor.extract_entities(text, **options)
|
||||
|
||||
# Extract relations if not provided
|
||||
if relations is None:
|
||||
@@ -373,15 +398,16 @@ class SemanticNetworkExtractor:
|
||||
# Pass method if specified
|
||||
if "relation_method" in self.config:
|
||||
rel_config["method"] = self.config["relation_method"]
|
||||
rel_extractor = RelationExtractor(
|
||||
**rel_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
relations = rel_extractor.extract_relations(text, entities, **options)
|
||||
if self._relation_extractor is None:
|
||||
self._relation_extractor = RelationExtractor(
|
||||
**rel_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
relations = self._relation_extractor.extract_relations(text, entities, **options)
|
||||
|
||||
# Build network
|
||||
total_steps = 2 # Create nodes, create edges
|
||||
|
||||
@@ -143,6 +143,13 @@ class TripletExtractor:
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
if method is not None:
|
||||
self.config["ner_method"] = method
|
||||
self.config["relation_method"] = method
|
||||
|
||||
self._ner_extractor = None
|
||||
self._relation_extractor = None
|
||||
|
||||
# Store parameters
|
||||
self.triplet_types = triplet_types
|
||||
self.include_temporal = include_temporal
|
||||
@@ -210,8 +217,12 @@ class TripletExtractor:
|
||||
message=f"Starting batch extraction... 0/{total_items}"
|
||||
)
|
||||
|
||||
# Determine max_workers
|
||||
max_workers = kwargs.get("max_workers", self.config.get("max_workers", 1))
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=self.method,
|
||||
)
|
||||
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
@@ -307,11 +318,12 @@ class TripletExtractor:
|
||||
|
||||
def extract_triplets(
|
||||
self,
|
||||
text: str,
|
||||
entities: Optional[List[Entity]] = None,
|
||||
relations: Optional[List[Relation]] = None,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> List[Triplet]:
|
||||
) -> Union[List[Triplet], List[List[Triplet]]]:
|
||||
"""
|
||||
Extract RDF triplets from text.
|
||||
|
||||
@@ -319,11 +331,29 @@ class TripletExtractor:
|
||||
text: Input text
|
||||
entities: Pre-extracted entities (optional)
|
||||
relations: Pre-extracted relations (optional)
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Extraction options
|
||||
|
||||
Returns:
|
||||
list: List of extracted triplets
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
entities_batch = entities
|
||||
if entities is not None and isinstance(entities, list) and (not entities or all(isinstance(e, Entity) for e in entities)):
|
||||
entities_batch = [entities for _ in range(len(text))] if entities else [[] for _ in range(len(text))]
|
||||
|
||||
relations_batch = relations
|
||||
if relations is not None and isinstance(relations, list) and (not relations or all(isinstance(r, Relation) for r in relations)):
|
||||
relations_batch = [relations for _ in range(len(text))] if relations else [[] for _ in range(len(text))]
|
||||
|
||||
return self.extract(
|
||||
text,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
pipeline_id=pipeline_id,
|
||||
**options,
|
||||
)
|
||||
|
||||
from .methods import get_triplet_method
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
@@ -341,16 +371,38 @@ class TripletExtractor:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Extracting entities..."
|
||||
)
|
||||
ner = NERExtractor(**self.config.get("ner", {}))
|
||||
entities = ner.extract_entities(text)
|
||||
if self._ner_extractor is None:
|
||||
ner_config = self.config.get("ner", {})
|
||||
if "ner_method" in self.config:
|
||||
ner_config = {**ner_config, "method": self.config["ner_method"]}
|
||||
self._ner_extractor = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation", "validator", "serializer", "quality"]
|
||||
},
|
||||
)
|
||||
entities = self._ner_extractor.extract_entities(text)
|
||||
|
||||
# Extract relations if not provided
|
||||
if relations is None:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Extracting relations..."
|
||||
)
|
||||
rel_extractor = RelationExtractor(**self.config.get("relation", {}))
|
||||
relations = rel_extractor.extract_relations(text, entities)
|
||||
if self._relation_extractor is None:
|
||||
rel_config = self.config.get("relation", {})
|
||||
if "relation_method" in self.config:
|
||||
rel_config = {**rel_config, "method": self.config["relation_method"]}
|
||||
self._relation_extractor = RelationExtractor(
|
||||
**rel_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation", "validator", "serializer", "quality"]
|
||||
},
|
||||
)
|
||||
relations = self._relation_extractor.extract_relations(text, entities)
|
||||
|
||||
# Use method-based extraction
|
||||
methods = options.get("method", self.method)
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
import statistics
|
||||
import time
|
||||
import os
|
||||
import sys
|
||||
from typing import Dict, List
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
from semantica.semantic_extract.event_detector import EventDetector
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.semantic_analyzer import SemanticAnalyzer
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.utils.progress_tracker import get_progress_tracker
|
||||
|
||||
|
||||
def _make_documents(n: int) -> List[Dict[str, str]]:
|
||||
base = (
|
||||
"Apple Inc. was founded by Steve Jobs in 1976 and is headquartered in Cupertino, California. "
|
||||
"Microsoft Corporation was founded by Bill Gates and Paul Allen in 1975. "
|
||||
"In 2014, Apple acquired Beats Electronics for $3 billion. "
|
||||
"In 2023, Google announced a partnership with OpenAI to improve search experiences."
|
||||
)
|
||||
return [{"id": f"doc_{i}", "content": f"{base} Document number {i}."} for i in range(n)]
|
||||
|
||||
|
||||
def _median_seconds(fn, repeats: int = 3) -> float:
|
||||
times = []
|
||||
for _ in range(repeats):
|
||||
start = time.perf_counter()
|
||||
fn()
|
||||
times.append(time.perf_counter() - start)
|
||||
return statistics.median(times)
|
||||
|
||||
|
||||
def _bench(label: str, fn, repeats: int = 3) -> dict:
|
||||
fn()
|
||||
seconds = _median_seconds(fn, repeats=repeats)
|
||||
return {"label": label, "seconds": seconds}
|
||||
|
||||
|
||||
def main():
|
||||
progress = get_progress_tracker()
|
||||
progress.displays = []
|
||||
|
||||
docs = _make_documents(80)
|
||||
texts = [d["content"] for d in docs]
|
||||
|
||||
ner = NERExtractor(method="pattern")
|
||||
rel = RelationExtractor(method="pattern")
|
||||
trip = TripletExtractor(method="pattern")
|
||||
events = EventDetector(method="pattern")
|
||||
analyzer = SemanticAnalyzer()
|
||||
net = SemanticNetworkExtractor(ner_method="pattern", relation_method="pattern")
|
||||
|
||||
results = []
|
||||
|
||||
def ner_parallel():
|
||||
ner.extract(texts)
|
||||
|
||||
def ner_seq():
|
||||
ner.extract(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("NER batch (default workers)", ner_parallel))
|
||||
results.append(_bench("NER batch (max_workers=1)", ner_seq))
|
||||
|
||||
entities_batch = ner.extract(texts)
|
||||
|
||||
def rel_parallel():
|
||||
rel.extract(texts, entities_batch)
|
||||
|
||||
def rel_seq():
|
||||
rel.extract(texts, entities_batch, max_workers=1)
|
||||
|
||||
results.append(_bench("Relation batch (default workers)", rel_parallel))
|
||||
results.append(_bench("Relation batch (max_workers=1)", rel_seq))
|
||||
|
||||
def trip_parallel():
|
||||
trip.extract(texts)
|
||||
|
||||
def trip_seq():
|
||||
trip.extract(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Triplet pipeline (default workers)", trip_parallel))
|
||||
results.append(_bench("Triplet pipeline (max_workers=1)", trip_seq))
|
||||
|
||||
def ev_parallel():
|
||||
events.detect_events(texts)
|
||||
|
||||
def ev_seq():
|
||||
events.detect_events(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Event detection (default workers)", ev_parallel))
|
||||
results.append(_bench("Event detection (max_workers=1)", ev_seq))
|
||||
|
||||
def analyzer_parallel():
|
||||
analyzer.analyze(texts)
|
||||
|
||||
def analyzer_seq():
|
||||
analyzer.analyze(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Semantic analysis (default workers)", analyzer_parallel))
|
||||
results.append(_bench("Semantic analysis (max_workers=1)", analyzer_seq))
|
||||
|
||||
def net_parallel():
|
||||
net.extract_network(texts)
|
||||
|
||||
def net_seq():
|
||||
net.extract_network(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Semantic network (default workers)", net_parallel))
|
||||
results.append(_bench("Semantic network (max_workers=1)", net_seq))
|
||||
|
||||
per_doc = []
|
||||
for row in results:
|
||||
per_doc.append({**row, "ms_per_doc": (row["seconds"] / len(texts)) * 1000.0})
|
||||
|
||||
print(f"Documents: {len(texts)}")
|
||||
for row in per_doc:
|
||||
print(f"{row['label']}: {row['seconds']:.3f}s ({row['ms_per_doc']:.2f} ms/doc)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -7,8 +7,12 @@ import os
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.ner_extractor import Entity as NEREntity
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.semantic_extract.event_detector import EventDetector
|
||||
from semantica.semantic_extract.semantic_analyzer import SemanticAnalyzer
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.named_entity_recognizer import Entity
|
||||
from semantica.semantic_extract.relation_extractor import Relation
|
||||
|
||||
@@ -51,6 +55,21 @@ class TestExtractors(unittest.TestCase):
|
||||
self.assertIsInstance(entities, list)
|
||||
mock_get_method.assert_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_entity_method")
|
||||
def test_ner_extraction_batch_via_extract_entities(self, mock_get_method):
|
||||
mock_method = MagicMock()
|
||||
mock_method.extract_entities.return_value = []
|
||||
mock_get_method.return_value = mock_method
|
||||
|
||||
extractor = NERExtractor(method="pattern")
|
||||
results = extractor.extract_entities(
|
||||
["Test text 1", "Test text 2"],
|
||||
)
|
||||
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_relation_method")
|
||||
def test_relation_extraction(self, mock_get_method):
|
||||
"""Test relation extraction call"""
|
||||
@@ -65,6 +84,21 @@ class TestExtractors(unittest.TestCase):
|
||||
self.assertIsInstance(relations, list)
|
||||
mock_get_method.assert_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_relation_method")
|
||||
def test_relation_extraction_batch_via_extract_relations(self, mock_get_method):
|
||||
mock_method = MagicMock()
|
||||
mock_method.extract_relations.return_value = []
|
||||
mock_get_method.return_value = mock_method
|
||||
|
||||
extractor = RelationExtractor(method="pattern")
|
||||
texts = ["A knows B", "A knows B"]
|
||||
entities = [NEREntity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)]
|
||||
|
||||
results = extractor.extract_relations(texts, entities)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_triplet_method")
|
||||
def test_triplet_extraction(self, mock_get_method):
|
||||
"""Test triplet extraction call"""
|
||||
@@ -81,5 +115,55 @@ class TestExtractors(unittest.TestCase):
|
||||
self.assertIsInstance(triplets, list)
|
||||
mock_get_method.assert_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_triplet_method")
|
||||
def test_triplet_extraction_batch_via_extract_triplets(self, mock_get_method):
|
||||
mock_method = MagicMock()
|
||||
mock_method.extract_triplets.return_value = []
|
||||
mock_get_method.return_value = mock_method
|
||||
|
||||
extractor = TripletExtractor(method="pattern")
|
||||
texts = ["A knows A", "A knows A"]
|
||||
entities_batch = [[NEREntity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)] for _ in texts]
|
||||
relations_batch = [[] for _ in texts]
|
||||
|
||||
results = extractor.extract_triplets(
|
||||
texts,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
)
|
||||
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
def test_event_detector_batch_via_detect_events(self):
|
||||
detector = EventDetector()
|
||||
texts = ["Apple acquired Beats in 2014.", "Google announced a partnership in 2023."]
|
||||
results = detector.detect_events(texts)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
def test_semantic_analyzer_batch_parallel(self):
|
||||
analyzer = SemanticAnalyzer()
|
||||
texts = ["A short sentence.", "Another short sentence."]
|
||||
results = analyzer.analyze(texts)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, dict) for r in results))
|
||||
|
||||
def test_semantic_network_batch_via_extract_network(self):
|
||||
extractor = SemanticNetworkExtractor()
|
||||
texts = ["A knows B.", "C knows D."]
|
||||
entities_batch = [[] for _ in texts]
|
||||
relations_batch = [[] for _ in texts]
|
||||
results = extractor.extract_network(
|
||||
texts,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
pytest.importorskip("groq")
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
|
||||
|
||||
def test_groq_llm_smoke_entities_relations_triplets():
|
||||
if not os.getenv("GROQ_API_KEY"):
|
||||
pytest.skip("GROQ_API_KEY is not set")
|
||||
|
||||
text = (
|
||||
"Apple acquired Beats in 2014 for $3 billion. "
|
||||
"Steve Jobs founded Apple. "
|
||||
"Beats is based in California."
|
||||
)
|
||||
model = "llama-3.3-70b-versatile"
|
||||
|
||||
entities = NERExtractor(method="llm").extract(
|
||||
text,
|
||||
provider="groq",
|
||||
model=model,
|
||||
temperature=0.0,
|
||||
max_tokens=250,
|
||||
)
|
||||
assert isinstance(entities, list)
|
||||
assert len(entities) > 0
|
||||
assert len(entities) <= 30
|
||||
|
||||
relations = RelationExtractor(method="llm").extract(
|
||||
text,
|
||||
entities=entities,
|
||||
provider="groq",
|
||||
model=model,
|
||||
temperature=0.0,
|
||||
max_tokens=350,
|
||||
max_entities_prompt=12,
|
||||
)
|
||||
assert isinstance(relations, list)
|
||||
assert len(relations) <= 30
|
||||
|
||||
triplets = TripletExtractor(method="llm").extract(
|
||||
text,
|
||||
entities=entities,
|
||||
relations=relations,
|
||||
provider="groq",
|
||||
model=model,
|
||||
temperature=0.0,
|
||||
max_tokens=350,
|
||||
)
|
||||
assert isinstance(triplets, list)
|
||||
assert len(triplets) <= 40
|
||||
@@ -168,12 +168,11 @@ class TestSemanticExtractImprovements(unittest.TestCase):
|
||||
texts = ["Text 1", "Text 2", "Text 3", "Text 4"]
|
||||
|
||||
start_time = time.time()
|
||||
# Run with 2 workers
|
||||
results = extractor.extract(texts, max_workers=2)
|
||||
results = extractor.extract(texts)
|
||||
end_time = time.time()
|
||||
|
||||
duration = end_time - start_time
|
||||
print(f" Parallel NER (2 workers) took {duration:.4f}s")
|
||||
print(f" Parallel NER (default workers) took {duration:.4f}s")
|
||||
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
@@ -206,11 +205,11 @@ class TestSemanticExtractImprovements(unittest.TestCase):
|
||||
entities = [[], [], [], []]
|
||||
|
||||
start_time = time.time()
|
||||
results = extractor.extract(texts, entities, max_workers=2)
|
||||
results = extractor.extract(texts, entities)
|
||||
end_time = time.time()
|
||||
|
||||
duration = end_time - start_time
|
||||
print(f" Parallel RE (2 workers) took {duration:.4f}s")
|
||||
print(f" Parallel RE (default workers) took {duration:.4f}s")
|
||||
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
@@ -264,11 +263,11 @@ class TestSemanticExtractImprovements(unittest.TestCase):
|
||||
texts = ["Text 1", "Text 2", "Text 3", "Text 4"]
|
||||
|
||||
start_time = time.time()
|
||||
results = extractor.extract(texts, max_workers=2)
|
||||
results = extractor.extract(texts)
|
||||
end_time = time.time()
|
||||
|
||||
duration = end_time - start_time
|
||||
print(f" Parallel TE (2 workers) took {duration:.4f}s")
|
||||
print(f" Parallel TE (default workers) took {duration:.4f}s")
|
||||
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
import multiprocessing
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from semantica.semantic_extract.config import resolve_max_workers
|
||||
from semantica.semantic_extract.ner_extractor import Entity, NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.methods import filter_entities_for_text
|
||||
from semantica.semantic_extract.schemas import RelationsResponse, RelationOut
|
||||
|
||||
|
||||
def test_resolve_max_workers_defaults_and_clamps():
|
||||
cpu_count = multiprocessing.cpu_count() or 1
|
||||
|
||||
assert resolve_max_workers(explicit=0) == 1
|
||||
assert resolve_max_workers(explicit=-10) == 1
|
||||
assert resolve_max_workers(explicit=1) == 1
|
||||
assert resolve_max_workers(explicit=10**9) == min(cpu_count, 32)
|
||||
|
||||
assert resolve_max_workers(explicit=None, methods=["ml"]) == 1
|
||||
|
||||
|
||||
def test_filter_entities_for_text_keeps_short_tokens():
|
||||
text = "US AI lab in NY"
|
||||
entities = [
|
||||
Entity(text="US", label="GPE", start_char=0, end_char=2, confidence=1.0),
|
||||
Entity(text="AI", label="TECH", start_char=3, end_char=5, confidence=1.0),
|
||||
Entity(text="NY", label="GPE", start_char=13, end_char=15, confidence=1.0),
|
||||
]
|
||||
kept = filter_entities_for_text(text, entities, max_keep=2)
|
||||
kept_texts = {e.text for e in kept}
|
||||
assert "US" in kept_texts or "AI" in kept_texts or "NY" in kept_texts
|
||||
|
||||
|
||||
def test_pattern_batch_defaults_to_single_worker_low_latency():
|
||||
extractor = NERExtractor(method="pattern")
|
||||
texts = [f"Text {i}" for i in range(8)]
|
||||
extractor.extract(texts)
|
||||
|
||||
|
||||
def test_relation_llm_prompt_filter_does_not_break_mapping():
|
||||
entities = [Entity(text=f"VeryLongEntityName{i}", label="ORG", start_char=0, end_char=1, confidence=1.0) for i in range(120)]
|
||||
ghost = Entity(text="Ghost", label="ORG", start_char=0, end_char=1, confidence=1.0)
|
||||
entities.append(ghost)
|
||||
|
||||
captured = {}
|
||||
|
||||
class FakeLLM:
|
||||
def is_available(self):
|
||||
return True
|
||||
|
||||
def generate_typed(self, prompt, schema, **kwargs):
|
||||
captured["prompt"] = prompt
|
||||
return RelationsResponse(
|
||||
relations=[
|
||||
RelationOut(subject="Ghost", predicate="related_to", object="VeryLongEntityName0", confidence=0.9)
|
||||
]
|
||||
)
|
||||
|
||||
with patch("semantica.semantic_extract.methods.create_provider", return_value=FakeLLM()):
|
||||
from semantica.semantic_extract.methods import extract_relations_llm
|
||||
|
||||
relations = extract_relations_llm(
|
||||
"Short text mentioning VeryLongEntityName0 only.",
|
||||
entities=entities,
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_entities_prompt=20,
|
||||
)
|
||||
|
||||
assert "Ghost" not in captured["prompt"]
|
||||
assert len(relations) == 1
|
||||
assert relations[0].subject.text == "Ghost"
|
||||
|
||||
|
||||
def test_triplet_extractor_reuses_sub_extractors():
|
||||
ner_instance = MagicMock()
|
||||
ner_instance.extract_entities.return_value = [
|
||||
Entity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)
|
||||
]
|
||||
|
||||
rel_instance = MagicMock()
|
||||
rel_instance.extract_relations.return_value = []
|
||||
|
||||
ner_ctor = MagicMock(return_value=ner_instance)
|
||||
rel_ctor = MagicMock(return_value=rel_instance)
|
||||
|
||||
with patch("semantica.semantic_extract.ner_extractor.NERExtractor", ner_ctor), patch(
|
||||
"semantica.semantic_extract.relation_extractor.RelationExtractor", rel_ctor
|
||||
), patch("semantica.semantic_extract.methods.get_triplet_method", return_value=lambda *args, **kwargs: []):
|
||||
extractor = TripletExtractor(method="pattern")
|
||||
extractor.extract_triplets("A text.")
|
||||
extractor.extract_triplets("A text again.")
|
||||
|
||||
assert ner_ctor.call_count == 1
|
||||
assert rel_ctor.call_count == 1
|
||||
|
||||
|
||||
def test_semantic_network_extractor_reuses_sub_extractors():
|
||||
ner_instance = MagicMock()
|
||||
ner_instance.extract_entities.return_value = [
|
||||
Entity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)
|
||||
]
|
||||
|
||||
rel_instance = MagicMock()
|
||||
rel_instance.extract_relations.return_value = []
|
||||
|
||||
ner_ctor = MagicMock(return_value=ner_instance)
|
||||
rel_ctor = MagicMock(return_value=rel_instance)
|
||||
|
||||
with patch("semantica.semantic_extract.ner_extractor.NERExtractor", ner_ctor), patch(
|
||||
"semantica.semantic_extract.relation_extractor.RelationExtractor", rel_ctor
|
||||
):
|
||||
extractor = SemanticNetworkExtractor(method="pattern")
|
||||
extractor.extract_network("A text.")
|
||||
extractor.extract_network("A text again.")
|
||||
|
||||
assert ner_ctor.call_count == 1
|
||||
assert rel_ctor.call_count == 1
|
||||
Reference in New Issue
Block a user