mseep-txtai 9.1.1__py3-none-any.whl

txtai/pipeline/llm/huggingface.py

@@ -0,0 +1,218 @@
+"""
+Hugging Face module
+"""
+
+from threading import Thread
+
+from transformers import AutoModelForImageTextToText, TextIteratorStreamer
+
+from ...models import Models
+
+from ..hfpipeline import HFPipeline
+
+from .generation import Generation
+
+
+class HFGeneration(Generation):
+    """
+    Hugging Face Transformers generative model.
+    """
+
+    def __init__(self, path, template=None, **kwargs):
+        # Call parent constructor
+        super().__init__(path, template, **kwargs)
+
+        # Create HuggingFace LLM pipeline
+        self.llm = HFLLM(path, **kwargs)
+
+    def isvision(self):
+        return isinstance(self.llm.pipeline.model, AutoModelForImageTextToText)
+
+    def stream(self, texts, maxlength, stream, stop, **kwargs):
+        yield from self.llm(texts, maxlength=maxlength, stream=stream, stop=stop, **kwargs)
+
+
+class HFLLM(HFPipeline):
+    """
+    Hugging Face Transformers large language model (LLM) pipeline. This pipeline autodetects if the model path
+    is a text generation or sequence to sequence model.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, task=None, **kwargs):
+        super().__init__(self.task(path, task, **kwargs), path, quantize, gpu, model, **kwargs)
+
+        # Load tokenizer, if necessary
+        self.pipeline.tokenizer = self.pipeline.tokenizer if self.pipeline.tokenizer else Models.tokenizer(path, **kwargs)
+
+    def __call__(self, text, prefix=None, maxlength=512, workers=0, stream=False, stop=None, **kwargs):
+        """
+        Generates text. Supports the following input formats:
+
+          - String or list of strings (instruction-tuned models must follow chat templates)
+          - List of dictionaries with `role` and `content` key-values or lists of lists
+
+        Args:
+            text: text|list
+            prefix: optional prefix to prepend to text elements
+            maxlength: maximum sequence length
+            workers: number of concurrent workers to use for processing data, defaults to None
+            stream: stream response if True, defaults to False
+            stop: list of stop strings
+            kwargs: additional generation keyword arguments
+
+        Returns:
+            generated text
+        """
+
+        # List of texts
+        texts = text if isinstance(text, list) else [text]
+
+        # Add prefix, if necessary
+        if prefix:
+            texts = [f"{prefix}{x}" for x in texts]
+
+        # Combine all keyword arguments
+        args, kwargs = self.parameters(texts, maxlength, workers, stop, **kwargs)
+
+        # Stream response
+        if stream:
+            return StreamingResponse(self.pipeline, texts, stop, **kwargs)()
+
+        # Run pipeline and extract generated text
+        results = [self.extract(result) for result in self.pipeline(*args, **kwargs)]
+
+        return results[0] if isinstance(text, str) else results
+
+    def parameters(self, texts, maxlength, workers, stop, **kwargs):
+        """
+        Builds a list of arguments and a combined parameter dictionary to use as keyword arguments.
+
+        Args:
+            texts: input texts
+            maxlength: maximum sequence length
+            workers: number of concurrent workers to use for processing data, defaults to None
+            stop: list of stop strings
+            kwargs: additional generation keyword arguments
+
+        Returns:
+            args, kwargs
+        """
+
+        # Set defaults and get underlying model
+        defaults, model = {"max_length": maxlength, "max_new_tokens": None, "num_workers": workers}, self.pipeline.model
+
+        # Set parameters for vision models and return
+        if self.pipeline.task == "image-text-to-text":
+            # Maxlength has to be large enough to accomodate images
+            defaults["max_length"] = max(maxlength, 2048)
+
+            # Set default token id
+            tokenid = model.generation_config.pad_token_id
+            model.generation_config.pad_token_id = tokenid if tokenid else model.generation_config.eos_token_id
+
+            # Vision models take all arguments as keyword arguments
+            return [], {**{"text": texts, "truncation": True}, **defaults, **kwargs}
+
+        # Add pad token if it's missing from model config
+        if not model.config.pad_token_id:
+            tokenid = model.config.eos_token_id
+            tokenid = tokenid[0] if isinstance(tokenid, list) else tokenid
+
+            # Set pad_token_id parameter
+            defaults["pad_token_id"] = tokenid
+
+            # Update tokenizer for batching
+            if "batch_size" in kwargs and self.pipeline.tokenizer.pad_token_id is None:
+                self.pipeline.tokenizer.pad_token_id = tokenid
+                self.pipeline.tokenizer.padding_side = "left"
+
+        # Set tokenizer when stop strings is set
+        if stop:
+            defaults["tokenizer"] = self.pipeline.tokenizer
+
+        return [texts], {**defaults, **kwargs}
+
+    def extract(self, result):
+        """
+        Extracts generated text from a pipeline result.
+
+        Args:
+            result: pipeline result
+
+        Returns:
+            generated text
+        """
+
+        # Extract output from list, if necessary
+        result = result[0] if isinstance(result, list) else result
+        text = result["generated_text"]
+        return text[-1]["content"] if isinstance(text, list) else text
+
+    def task(self, path, task, **kwargs):
+        """
+        Get the pipeline task name.
+
+        Args:
+            path: model path input
+            task: task name
+            kwargs: optional additional keyword arguments
+
+        Returns:
+            pipeline task name
+        """
+
+        # Mapping from txtai to Hugging Face pipeline tasks
+        mapping = {"language-generation": "text-generation", "sequence-sequence": "text2text-generation", "vision": "image-text-to-text"}
+
+        # Attempt to resolve task
+        if path and not task:
+            task = Models.task(path, **kwargs)
+
+        # Map to Hugging Face task. Default to text2text-generation pipeline when task not resolved.
+        return mapping.get(task, "text2text-generation")
+
+
+class Generator(HFLLM):
+    """
+    Generate text with a causal language model.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs):
+        super().__init__(path, quantize, gpu, model, "language-generation", **kwargs)
+
+
+class Sequences(HFLLM):
+    """
+    Generate text with a sequence-sequence model.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs):
+        super().__init__(path, quantize, gpu, model, "sequence-sequence", **kwargs)
+
+
+class StreamingResponse:
+    """
+    Generate text as a streaming response.
+    """
+
+    def __init__(self, pipeline, texts, stop, **kwargs):
+        # Create streamer
+        self.stream = TextIteratorStreamer(pipeline.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=5)
+        kwargs["streamer"] = self.stream
+        kwargs["stop_strings"] = stop
+
+        # Create thread
+        self.thread = Thread(target=pipeline, args=[texts], kwargs=kwargs)
+
+        # Store number of inputs
+        self.length = len(texts)
+
+    def __call__(self):
+        # Start the process
+        self.thread.start()
+
+        return self
+
+    def __iter__(self):
+        for _ in range(self.length):
+            yield from self.stream

txtai/pipeline/llm/litellm.py

@@ -0,0 +1,90 @@
+"""
+LiteLLM module
+"""
+
+from transformers.utils import cached_file
+
+# Conditional import
+try:
+    import litellm as api
+
+    LITELLM = True
+except ImportError:
+    LITELLM = False
+
+from .generation import Generation
+
+
+class LiteLLM(Generation):
+    """
+    LiteLLM generative model.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is a LiteLLM model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is a LiteLLM model, False otherwise
+        """
+
+        # pylint: disable=W0702
+        if isinstance(path, str) and LITELLM:
+            debug = api.suppress_debug_info
+            try:
+                # Suppress debug messages for this test
+                api.suppress_debug_info = True
+                return api.get_llm_provider(path) and not LiteLLM.ishub(path)
+            except:
+                return False
+            finally:
+                # Restore debug info value to original value
+                api.suppress_debug_info = debug
+
+        return False
+
+    @staticmethod
+    def ishub(path):
+        """
+        Checks if path is available on the HF Hub.
+
+        Args:
+            input path
+
+        Returns:
+            True if this is a model on the HF Hub
+        """
+
+        # pylint: disable=W0702
+        try:
+            return cached_file(path_or_repo_id=path, filename="config.json") is not None if "/" in path else False
+        except:
+            return False
+
+    def __init__(self, path, template=None, **kwargs):
+        super().__init__(path, template, **kwargs)
+
+        if not LITELLM:
+            raise ImportError('LiteLLM is not available - install "pipeline" extra to enable')
+
+        # Ignore common pipeline parameters
+        self.kwargs = {k: v for k, v in self.kwargs.items() if k not in ["quantize", "gpu", "model", "task"]}
+
+    def stream(self, texts, maxlength, stream, stop, **kwargs):
+        for text in texts:
+            # LLM API call
+            result = api.completion(
+                model=self.path,
+                messages=[{"content": text, "role": "prompt"}] if isinstance(text, str) else text,
+                max_tokens=maxlength,
+                stream=stream,
+                stop=stop,
+                **{**self.kwargs, **kwargs}
+            )
+
+            # Stream response
+            yield from self.response(result if stream else [result])

txtai/pipeline/llm/llama.py

@@ -0,0 +1,152 @@
+"""
+Llama module
+"""
+
+import os
+
+from huggingface_hub import hf_hub_download
+
+# Conditional import
+try:
+    import llama_cpp as llama
+
+    LLAMA_CPP = True
+except ImportError:
+    LLAMA_CPP = False
+
+from .generation import Generation
+
+
+class LlamaCpp(Generation):
+    """
+    llama.cpp generative model.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is a llama.cpp model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is a llama.cpp model, False otherwise
+        """
+
+        return isinstance(path, str) and path.lower().endswith(".gguf")
+
+    def __init__(self, path, template=None, **kwargs):
+        super().__init__(path, template, **kwargs)
+
+        if not LLAMA_CPP:
+            raise ImportError('llama.cpp is not available - install "pipeline" extra to enable')
+
+        # Check if this is a local path, otherwise download from the HF Hub
+        path = path if os.path.exists(path) else self.download(path)
+
+        # Create llama.cpp instance
+        self.llm = self.create(path, **kwargs)
+
+    def stream(self, texts, maxlength, stream, stop, **kwargs):
+        for text in texts:
+            yield from (
+                self.messages(text, maxlength, stream, stop, **kwargs)
+                if isinstance(text, list)
+                else self.prompt(text, maxlength, stream, stop, **kwargs)
+            )
+
+    def download(self, path):
+        """
+        Downloads path from the Hugging Face Hub.
+
+        Args:
+            path: full model path
+
+        Returns:
+            local cached model path
+        """
+
+        # Split into parts
+        parts = path.split("/")
+
+        # Calculate repo id split
+        repo = 2 if len(parts) > 2 else 1
+
+        # Download and cache file
+        return hf_hub_download(repo_id="/".join(parts[:repo]), filename="/".join(parts[repo:]))
+
+    def create(self, path, **kwargs):
+        """
+        Creates a new llama.cpp model instance.
+
+        Args:
+            path: path to model
+            kwargs: additional keyword args
+
+        Returns:
+            llama.cpp instance
+        """
+
+        # Default n_ctx=0 if not already set. This sets n_ctx = n_ctx_train.
+        kwargs["n_ctx"] = kwargs.get("n_ctx", 0)
+
+        # Default GPU layers if not already set
+        kwargs["n_gpu_layers"] = kwargs.get("n_gpu_layers", -1 if kwargs.get("gpu", os.environ.get("LLAMA_NO_METAL") != "1") else 0)
+
+        # Default verbose flag
+        kwargs["verbose"] = kwargs.get("verbose", False)
+
+        # Create llama.cpp instance
+        try:
+            return llama.Llama(model_path=path, **kwargs)
+        except ValueError as e:
+            # Fallback to default n_ctx when not enough memory for n_ctx = n_ctx_train
+            if not kwargs["n_ctx"]:
+                kwargs.pop("n_ctx")
+                return llama.Llama(model_path=path, **kwargs)
+
+            # Raise exception if n_ctx manually specified
+            raise e
+
+    def messages(self, messages, maxlength, stream, stop, **kwargs):
+        """
+        Processes a list of messages.
+
+        Args:
+            messages: list of dictionaries with `role` and `content` key-values
+            maxlength: maximum sequence length
+            stream: stream response if True, defaults to False
+            stop: list of stop strings
+            kwargs: additional generation keyword arguments
+
+        Returns:
+            generated text
+        """
+
+        # LLM call with messages
+        result = self.llm.create_chat_completion(messages=messages, max_tokens=maxlength, stream=stream, stop=stop, **kwargs)
+
+        # Stream response
+        yield from self.response(result if stream else [result])
+
+    def prompt(self, text, maxlength, stream, stop, **kwargs):
+        """
+        Processes a prompt.
+
+        Args:
+            prompt: prompt text
+            maxlength: maximum sequence length
+            stream: stream response if True, defaults to False
+            stop: list of stop strings
+            kwargs: additional generation keyword arguments
+
+        Returns:
+            generated text
+        """
+
+        # LLM call with prompt
+        result = self.llm(text, max_tokens=maxlength, stream=stream, stop=stop, **kwargs)
+
+        # Stream response
+        yield from self.response(result if stream else [result])

txtai/pipeline/llm/llm.py

@@ -0,0 +1,75 @@
+"""
+LLM module
+"""
+
+import logging
+
+from .factory import GenerationFactory
+
+from ..base import Pipeline
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+
+class LLM(Pipeline):
+    """
+    Pipeline for running large language models (LLMs). This class supports the following LLM backends:
+
+      - Local LLMs with Hugging Face Transformers
+      - Local LLMs with llama.cpp
+      - Remote API LLMs with LiteLLM
+      - Custom generation implementations
+    """
+
+    def __init__(self, path=None, method=None, **kwargs):
+        """
+        Creates a new LLM.
+
+        Args:
+            path: model path
+            method: llm model framework, infers from path if not provided
+            kwargs: model keyword arguments
+        """
+
+        # Default LLM if not provided
+        path = path if path else "google/flan-t5-base"
+
+        # Generation instance
+        self.generator = GenerationFactory.create(path, method, **kwargs)
+
+    def __call__(self, text, maxlength=512, stream=False, stop=None, defaultrole="prompt", stripthink=False, **kwargs):
+        """
+        Generates text. Supports the following input formats:
+
+          - String or list of strings (instruction-tuned models must follow chat templates)
+          - List of dictionaries with `role` and `content` key-values or lists of lists
+
+        Args:
+            text: text|list
+            maxlength: maximum sequence length
+            stream: stream response if True, defaults to False
+            stop: list of stop strings, defaults to None
+            defaultrole: default role to apply to text inputs (prompt for raw prompts (default) or user for user chat messages)
+            stripthink: strip thinking tags, defaults to False
+            kwargs: additional generation keyword arguments
+
+        Returns:
+            generated text
+        """
+
+        # Debug logging
+        logger.debug(text)
+
+        # Run LLM generation
+        return self.generator(text, maxlength, stream, stop, defaultrole, stripthink, **kwargs)
+
+    def isvision(self):
+        """
+        Returns True if this LLM supports vision operations.
+
+        Returns:
+            True if this is a vision model
+        """
+
+        return self.generator.isvision()