"""NuExtract 1.5 Smol adapter for on-demand CPU extraction. Provides an isolated interface for NuExtract inference: - Production mode: loads numind/NuExtract-1.5-smol on CPU - Test mode: deterministic schema-based mock extraction This adapter uses the shared InferenceGateway with a configurable target for CPU-only deployment. It is NOT always-resident — loaded on demand for benchmark or promoted document classes only. Requirement: 6.6 """ from __future__ import annotations import logging import re import time from typing import Any from services.intelligence_pipeline_v3.nuextract.models import ( ExtractedField, NuExtractResult, ) logger = logging.getLogger(__name__) # Pinned model configuration NUEXTRACT_MODEL_NAME = "numind/NuExtract-1.5-smol" NUEXTRACT_MODEL_VERSION = "numind/NuExtract-1.5-smol@v1.5" class NuExtractAdapter: """Adapter for NuExtract 1.5 Smol hierarchical extraction. Designed for CPU/on-demand use, not always-resident GPU deployment. Parameters ---------- test_mode When True, uses deterministic schema-based extraction rather than loading the model. Useful for testing without model dependencies. max_length Maximum input text length in characters. Longer texts are processed in segments. """ def __init__( self, test_mode: bool = True, max_length: int = 16_000, ) -> None: self._test_mode = test_mode self._max_length = max_length self._model = None self._tokenizer = None self._model_name = NUEXTRACT_MODEL_NAME self._model_version = NUEXTRACT_MODEL_VERSION self._loaded = False if not test_mode: self._load_model() @property def model_version(self) -> str: """Return the pinned model version string.""" return self._model_version @property def model_name(self) -> str: """Return the model name.""" return self._model_name @property def is_loaded(self) -> bool: """Return whether the model is currently loaded.""" return self._loaded def _load_model(self) -> None: """Load the NuExtract model and tokenizer for CPU inference.""" try: from transformers import AutoModelForCausalLM, AutoTokenizer logger.info("Loading NuExtract model (CPU): %s", self._model_name) self._tokenizer = AutoTokenizer.from_pretrained(self._model_name) self._model = AutoModelForCausalLM.from_pretrained( self._model_name, device_map="cpu", torch_dtype="auto", ) self._model.eval() self._loaded = True logger.info("NuExtract model loaded successfully on CPU") except ImportError: raise RuntimeError( "transformers and torch are required for NuExtract inference. " "Install with: pip install transformers torch" ) except Exception as e: raise RuntimeError(f"Failed to load NuExtract model: {e}") from e def unload(self) -> None: """Unload model to free memory (on-demand lifecycle).""" if self._model is not None: del self._model del self._tokenizer self._model = None self._tokenizer = None self._loaded = False logger.info("NuExtract model unloaded") async def extract( self, text: str, schema: dict[str, Any], document_type: str = "", ) -> NuExtractResult: """Extract structured fields from text using the given schema. Parameters ---------- text Source document text. schema JSON schema defining the fields to extract. document_type Document type identifier for reporting. Returns ------- NuExtractResult Extracted fields with confidence, latency, and memory metrics. """ start_time = time.perf_counter() try: if self._test_mode: fields = self._extract_test_mode(text, schema) else: fields = self._extract_production(text, schema) latency_ms = (time.perf_counter() - start_time) * 1000 memory_mb = self._estimate_memory() # Compute overall confidence as mean of field confidences confidence = 0.0 if fields: confidence = sum(f.confidence for f in fields) / len(fields) return NuExtractResult( fields=fields, spans=[ {"start": f.start_char, "end": f.end_char, "field": f.name} for f in fields if f.start_char is not None ], confidence=confidence, model_version=self._model_version, latency_ms=latency_ms, memory_mb=memory_mb, document_type=document_type, schema_used=schema, ) except Exception as e: latency_ms = (time.perf_counter() - start_time) * 1000 logger.error("NuExtract extraction failed: %s", e) return NuExtractResult( model_version=self._model_version, latency_ms=latency_ms, document_type=document_type, schema_used=schema, error=str(e), ) def _extract_test_mode( self, text: str, schema: dict[str, Any] ) -> list[ExtractedField]: """Deterministic extraction for testing. Matches schema field names against text content using simple pattern matching to simulate extraction behavior. """ fields: list[ExtractedField] = [] properties = schema.get("properties", schema) for field_name, field_spec in properties.items(): # Simple pattern: look for the field name or related keywords in text pattern = re.compile( rf"\b{re.escape(field_name.replace('_', ' '))}[:\s]+([^\n.;]+)", re.IGNORECASE, ) match = pattern.search(text) if match: value = match.group(1).strip() fields.append( ExtractedField( name=field_name, value=value, start_char=match.start(1), end_char=match.end(1), confidence=0.85, ) ) else: # Try extracting from nearby context for hierarchical schemas if isinstance(field_spec, dict) and "properties" in field_spec: # Nested schema — attempt hierarchical extraction nested = self._extract_test_mode(text, field_spec) if nested: fields.append( ExtractedField( name=field_name, value={f.name: f.value for f in nested}, confidence=sum(f.confidence for f in nested) / len(nested), ) ) else: # Field not found in text fields.append( ExtractedField( name=field_name, value=None, confidence=0.0, ) ) return fields def _extract_production( self, text: str, schema: dict[str, Any] ) -> list[ExtractedField]: """Run NuExtract inference on text using the loaded model.""" import json import torch if self._model is None or self._tokenizer is None: raise RuntimeError("Model not loaded. Initialize with test_mode=False.") # Format input in NuExtract's expected format schema_str = json.dumps(schema, indent=2) prompt = f"<|input|>\n### Template:\n{schema_str}\n### Text:\n{text[:self._max_length]}\n<|output|>\n" inputs = self._tokenizer( prompt, return_tensors="pt", truncation=True, max_length=4096, ) with torch.no_grad(): outputs = self._model.generate( **inputs, max_new_tokens=1024, temperature=0.0, do_sample=False, ) # Decode and parse the output generated = self._tokenizer.decode( outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True, ) return self._parse_output(generated, text, schema) def _parse_output( self, output: str, source_text: str, schema: dict[str, Any] ) -> list[ExtractedField]: """Parse model output into structured fields with spans.""" import json fields: list[ExtractedField] = [] try: parsed = json.loads(output) except json.JSONDecodeError: logger.warning("Failed to parse NuExtract output as JSON") return fields properties = schema.get("properties", schema) for field_name in properties: if field_name in parsed: value = parsed[field_name] # Try to find the value in source text for span start_char = None end_char = None if isinstance(value, str) and value: idx = source_text.find(value) if idx >= 0: start_char = idx end_char = idx + len(value) fields.append( ExtractedField( name=field_name, value=value, start_char=start_char, end_char=end_char, confidence=0.8, ) ) return fields def _estimate_memory(self) -> float: """Estimate current memory usage in MB.""" if self._test_mode: return 0.0 try: import torch if torch.cuda.is_available(): return torch.cuda.memory_allocated() / (1024 * 1024) # For CPU, estimate from model parameters if self._model is not None: param_bytes = sum( p.nelement() * p.element_size() for p in self._model.parameters() ) return param_bytes / (1024 * 1024) except Exception: pass return 0.0