Source code for langchain_experimental.llms.lmformatenforcer_decoder
"""Experimental implementation of lm-format-enforcer wrapped LLM."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, List, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from langchain.schema import LLMResult
from langchain_experimental.pydantic_v1 import Field
if TYPE_CHECKING:
    import lmformatenforcer
[docs]def import_lmformatenforcer() -> lmformatenforcer:
    """Lazily import lmformatenforcer."""
    try:
        import lmformatenforcer
    except ImportError:
        raise ImportError(
            "Could not import lmformatenforcer python package. "
            "Please install it with `pip install lm-format-enforcer`."
        )
    return lmformatenforcer
[docs]class LMFormatEnforcer(HuggingFacePipeline):
    """LMFormatEnforcer wrapped LLM using HuggingFace Pipeline API.
    This pipeline is experimental and not yet stable.
    """
    json_schema: Optional[dict] = Field(
        description="The JSON Schema to complete.", default=None
    )
    regex: Optional[str] = Field(
        description="The regular expression to complete.", default=None
    )
    def _generate(
        self,
        prompts: List[str],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> LLMResult:
        lmformatenforcer = import_lmformatenforcer()
        import lmformatenforcer.integrations.transformers as hf_integration