codekingpro/portable-devtools

1from __future__ import annotations2 3import importlib.util4import logging5import platform6from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Union7 8from langchain_core.callbacks import CallbackManagerForLLMRun9from langchain_core.language_models.llms import LLM10from langchain_core.outputs import GenerationChunk11from pydantic import BaseModel, Field, model_validator12 13logger = logging.getLogger(__name__)14 15 16class Outlines(LLM):17    """LLM wrapper for the Outlines library."""18 19    client: Any = None  # :meta private:20 21    model: str22    """Identifier for the model to use with Outlines.23    24    The model identifier should be a string specifying:25    - A Hugging Face model name (e.g., "meta-llama/Llama-2-7b-chat-hf")26    - A local path to a model27    - For GGUF models, the format is "repo_id/file_name"28      (e.g., "TheBloke/Llama-2-7B-Chat-GGUF/llama-2-7b-chat.Q4_K_M.gguf")29    30    Examples:31    - "TheBloke/Llama-2-7B-Chat-GGUF/llama-2-7b-chat.Q4_K_M.gguf"32    - "meta-llama/Llama-2-7b-chat-hf"33    """34 35    backend: Literal[36        "llamacpp", "transformers", "transformers_vision", "vllm", "mlxlm"37    ] = "transformers"38    """Specifies the backend to use for the model.39    40    Supported backends are:41    - "llamacpp": For GGUF models using llama.cpp42    - "transformers": For Hugging Face Transformers models (default)43    - "transformers_vision": For vision-language models (e.g., LLaVA)44    - "vllm": For models using the vLLM library45    - "mlxlm": For models using the MLX framework46    47    Note: Ensure you have the necessary dependencies installed for the chosen backend.48    The system will attempt to import required packages and may raise an ImportError49    if they are not available.50    """51 52    max_tokens: int = 25653    """The maximum number of tokens to generate."""54 55    stop: Optional[List[str]] = None56    """A list of strings to stop generation when encountered."""57 58    streaming: bool = True59    """Whether to stream the results, token by token."""60 61    regex: Optional[str] = None62    r"""Regular expression for structured generation.63    64    If provided, Outlines will guarantee that the generated text matches this regex.65    This can be useful for generating structured outputs like IP addresses, dates, etc.66    67    Example: (valid IP address)68        regex = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"69    70    Note: Computing the regex index can take some time, so it's recommended to reuse71    the same regex for multiple generations if possible.72    73    For more details, see: https://dottxt-ai.github.io/outlines/reference/generation/regex/74    """75 76    type_constraints: Optional[Union[type, str]] = None77    """Type constraints for structured generation.78    79    Restricts the output to valid Python types. Supported types include:80    int, float, bool, datetime.date, datetime.time, datetime.datetime.81    82    Example:83        type_constraints = int84    85    For more details, see: https://dottxt-ai.github.io/outlines/reference/generation/format/86    """87 88    json_schema: Optional[Union[BaseModel, Dict, Callable]] = None89    """Pydantic model, JSON Schema, or callable (function signature)90    for structured JSON generation.91    92    Outlines can generate JSON output that follows a specified structure,93    which is useful for:94    1. Parsing the answer (e.g., with Pydantic), storing it, or returning it to a user.95    2. Calling a function with the result.96 97    You can provide:98    - A Pydantic model99    - A JSON Schema (as a Dict)100    - A callable (function signature)101 102    The generated JSON will adhere to the specified structure.103 104    For more details, see: https://dottxt-ai.github.io/outlines/reference/generation/json/105    """106 107    grammar: Optional[str] = None108    """Context-free grammar for structured generation.109    110    If provided, Outlines will generate text that adheres to the specified grammar.111    The grammar should be defined in EBNF format.112    113    This can be useful for generating structured outputs like mathematical expressions,114    programming languages, or custom domain-specific languages.115    116    Example:117        grammar = '''118            ?start: expression119            ?expression: term (("+" | "-") term)*120            ?term: factor (("*" | "/") factor)*121            ?factor: NUMBER | "-" factor | "(" expression ")"122            %import common.NUMBER123        '''124    125    Note: Grammar-based generation is currently experimental and may have performance126    limitations. It uses greedy generation to mitigate these issues.127    128    For more details and examples, see:129    https://dottxt-ai.github.io/outlines/reference/generation/cfg/130    """131 132    custom_generator: Optional[Any] = None133    """Set your own outlines generator object to override the default behavior."""134 135    model_kwargs: Dict[str, Any] = Field(default_factory=dict)136    """Additional parameters to pass to the underlying model.137    138    Example:139        model_kwargs = {"temperature": 0.8, "seed": 42}140    """141 142    @model_validator(mode="after")143    def validate_environment(self) -> "Outlines":144        """Validate that outlines is installed and create a model instance."""145        num_constraints = sum(146            [147                bool(self.regex),148                bool(self.type_constraints),149                bool(self.json_schema),150                bool(self.grammar),151            ]152        )153        if num_constraints > 1:154            raise ValueError(155                "Either none or exactly one of regex, type_constraints, "156                "json_schema, or grammar can be provided."157            )158        return self.build_client()159 160    def build_client(self) -> "Outlines":161        try:162            import outlines.models as models163        except ImportError:164            raise ImportError(165                "Could not import the Outlines library. "166                "Please install it with `pip install outlines`."167            )168 169        def check_packages_installed(170            packages: List[Union[str, Tuple[str, str]]],171        ) -> None:172            missing_packages = [173                pkg if isinstance(pkg, str) else pkg[0]174                for pkg in packages175                if importlib.util.find_spec(pkg[1] if isinstance(pkg, tuple) else pkg)176                is None177            ]178            if missing_packages:179                raise ImportError(  # todo this is displaying wrong180                    f"Missing packages: {', '.join(missing_packages)}. "181                    "You can install them with:\n\n"182                    f"    pip install {' '.join(missing_packages)}"183                )184 185        if self.backend == "llamacpp":186            if ".gguf" in self.model:187                creator, repo_name, file_name = self.model.split("/", 2)188                repo_id = f"{creator}/{repo_name}"189            else:  # todo add auto-file-selection if no file is given190                raise ValueError("GGUF file_name must be provided for llama.cpp.")191            check_packages_installed([("llama-cpp-python", "llama_cpp")])192            self.client = models.llamacpp(repo_id, file_name, **self.model_kwargs)193        elif self.backend == "transformers":194            check_packages_installed(["transformers", "torch", "datasets"])195            self.client = models.transformers(self.model, **self.model_kwargs)196        elif self.backend == "transformers_vision":197            check_packages_installed(198                [199                    "transformers",200                    "datasets",201                    "torchvision",202                    "PIL",203                    "flash_attn",204                ]205            )206            from transformers import LlavaNextForConditionalGeneration207 208            if not hasattr(models, "transformers_vision"):209                raise ValueError(210                    "transformers_vision backend is not supported, "211                    "please install the correct outlines version."212                )213            self.client = models.transformers_vision(214                self.model,215                model_class=LlavaNextForConditionalGeneration,216                **self.model_kwargs,217            )218        elif self.backend == "vllm":219            if platform.system() == "Darwin":220                raise ValueError("vLLM backend is not supported on macOS.")221            check_packages_installed(["vllm"])222            self.client = models.vllm(self.model, **self.model_kwargs)223        elif self.backend == "mlxlm":224            check_packages_installed(["mlx"])225            self.client = models.mlxlm(self.model, **self.model_kwargs)226        else:227            raise ValueError(f"Unsupported backend: {self.backend}")228 229        return self230 231    @property232    def _llm_type(self) -> str:233        return "outlines"234 235    @property236    def _default_params(self) -> Dict[str, Any]:237        return {238            "max_tokens": self.max_tokens,239            "stop_at": self.stop,240            **self.model_kwargs,241        }242 243    @property244    def _identifying_params(self) -> Dict[str, Any]:245        return {246            "model": self.model,247            "backend": self.backend,248            "regex": self.regex,249            "type_constraints": self.type_constraints,250            "json_schema": self.json_schema,251            "grammar": self.grammar,252            **self._default_params,253        }254 255    @property256    def _generator(self) -> Any:257        from outlines import generate258 259        if self.custom_generator:260            return self.custom_generator261        if self.regex:262            return generate.regex(self.client, regex_str=self.regex)263        if self.type_constraints:264            return generate.format(self.client, python_type=self.type_constraints)265        if self.json_schema:266            return generate.json(self.client, schema_object=self.json_schema)267        if self.grammar:268            return generate.cfg(self.client, cfg_str=self.grammar)269        return generate.text(self.client)270 271    def _call(272        self,273        prompt: str,274        stop: Optional[List[str]] = None,275        run_manager: Optional[CallbackManagerForLLMRun] = None,276        **kwargs: Any,277    ) -> str:278        params = {**self._default_params, **kwargs}279        if stop:280            params["stop_at"] = stop281 282        response = ""283        if self.streaming:284            for chunk in self._stream(285                prompt=prompt,286                stop=params["stop_at"],287                run_manager=run_manager,288                **params,289            ):290                response += chunk.text291        else:292            response = self._generator(prompt, **params)293        return response294 295    def _stream(296        self,297        prompt: str,298        stop: Optional[List[str]] = None,299        run_manager: Optional[CallbackManagerForLLMRun] = None,300        **kwargs: Any,301    ) -> Iterator[GenerationChunk]:302        params = {**self._default_params, **kwargs}303        if stop:304            params["stop_at"] = stop305 306        for token in self._generator.stream(prompt, **params):307            if run_manager:308                run_manager.on_llm_new_token(token)309            yield GenerationChunk(text=token)310 311    @property312    def tokenizer(self) -> Any:313        """Access the tokenizer for the underlying model.314 315        .encode() to tokenize text.316        .decode() to convert tokens back to text.317        """318        if hasattr(self.client, "tokenizer"):319            return self.client.tokenizer320        raise ValueError("Tokenizer not found")321