diff --git a/gliner2/inference/engine.py b/gliner2/inference/engine.py index 08a9062..730aa6c 100644 --- a/gliner2/inference/engine.py +++ b/gliner2/inference/engine.py @@ -31,7 +31,7 @@ import json from collections import OrderedDict from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Union, Tuple, TYPE_CHECKING, Pattern, Literal +from typing import Any, Dict, List, Optional, Union, Tuple, TYPE_CHECKING, Pattern, Literal, Type, TypeVar, get_origin, get_args import torch import torch.nn.functional as F @@ -43,6 +43,102 @@ if TYPE_CHECKING: from gliner2.api_client import GLiNER2API +try: + from pydantic import BaseModel + PYDANTIC_AVAILABLE = True +except ImportError: + BaseModel = Any + PYDANTIC_AVAILABLE = False + +T = TypeVar("T", bound="BaseModel") + + +# ============================================================================= +# Pydantic Integration Utilities +# ============================================================================= + +class PydanticConverter: + """Bridge for Pydantic-based schema generation and object materialization.""" + + @staticmethod + def to_schema(model_class: Type[BaseModel], schema_builder: 'Schema') -> 'Schema': + """Generates a GLiNER2 structure schema from a Pydantic model.""" + if not PYDANTIC_AVAILABLE: + raise ImportError("Pydantic is required for this feature. Install it with `pip install pydantic`.") + + struct_name = model_class.__name__.lower() + builder = schema_builder.structure(struct_name) + + fields = getattr(model_class, "model_fields", getattr(model_class, "__fields__", {})) + + for name, field in fields.items(): + annotation = getattr(field, "annotation", getattr(field, "type_", None)) + origin = get_origin(annotation) + args = get_args(annotation) + + # Prevent silent failures by warning about unsupported nested models + check_type = annotation + if origin: + check_type = args[0] if args else annotation + + try: + if isinstance(check_type, type) and issubclass(check_type, BaseModel): + import warnings + warnings.warn( + f"Field '{name}' in {model_class.__name__} appears to be a nested Pydantic model. " + "GLiNER2 currently only supports flat structures; this field will be extracted as raw text.", + UserWarning + ) + except TypeError: + pass + + dtype: Literal["str", "list"] = "list" if origin is list else "str" + choices = None + + if origin is Literal: + choices = list(args) + dtype = "str" + elif origin is Union: + non_none_args = [a for a in args if a is not type(None)] + if non_none_args: + inner_origin = get_origin(non_none_args[0]) + dtype = "list" if inner_origin is list else "str" + if inner_origin is Literal: + choices = list(get_args(non_none_args[0])) + dtype = "str" + + description = getattr(field, "description", None) + + builder.field( + name=name, + dtype=dtype, + choices=choices, + description=description + ) + + return schema_builder + + @staticmethod + def materialize(model_class: Type[T], data: Dict[str, Any]) -> Optional[T]: + """Instantiates a Pydantic model from extracted dictionary data.""" + if not PYDANTIC_AVAILABLE: + raise ImportError("Pydantic is required for this feature. Install it with `pip install pydantic`.") + + struct_name = model_class.__name__.lower() + instances = data.get(struct_name, []) + if not instances: + return None + + # Map internal multi-instance output to single model materialization + instance_data = instances[0] + + # Exclude nulls to allow Pydantic to apply field defaults and handle Optional types + cleaned_data = {k: v for k, v in instance_data.items() if v is not None} + + if hasattr(model_class, "model_validate"): + return model_class.model_validate(cleaned_data) + return model_class(**cleaned_data) + # ============================================================================= # Validators @@ -288,6 +384,10 @@ def build(self) -> Dict[str, Any]: self._active_builder = None return self.schema + def from_pydantic(self, model_class: Type['BaseModel']) -> 'Schema': + """Builds schema from a Pydantic model class.""" + return PydanticConverter.to_schema(model_class, self) + # ============================================================================= # Main GLiNER2 Class @@ -316,6 +416,29 @@ def create_schema(self) -> Schema: """Create a new schema builder.""" return Schema() + def predict_model( + self, + text: str, + model_class: Type[T], + threshold: float = 0.5, + **kwargs + ) -> Optional[T]: + """ + Performs extraction and materializes the result into a Pydantic model. + + Args: + text: Input text to extract from + model_class: Pydantic model class defining the schema + threshold: Confidence threshold for extraction + **kwargs: Additional arguments for extract() + + Returns: + Materialized Pydantic model instance or None. + """ + schema = self.create_schema().from_pydantic(model_class) + raw_result = self.extract(text, schema, threshold=threshold, **kwargs) + return PydanticConverter.materialize(model_class, raw_result) + # ========================================================================= # Main Batch Extraction # =========================================================================