mseep-txtai 9.1.1__py3-none-any.whl

txtai/pipeline/image/caption.py

@@ -0,0 +1,55 @@
+"""
+Caption module
+"""
+
+# Conditional import
+try:
+    from PIL import Image
+
+    PIL = True
+except ImportError:
+    PIL = False
+
+from ..hfpipeline import HFPipeline
+
+
+class Caption(HFPipeline):
+    """
+    Constructs captions for images.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs):
+        if not PIL:
+            raise ImportError('Captions pipeline is not available - install "pipeline" extra to enable')
+
+        # Call parent constructor
+        super().__init__("image-to-text", path, quantize, gpu, model, **kwargs)
+
+    def __call__(self, images):
+        """
+        Builds captions for images.
+
+        This method supports a single image or a list of images. If the input is an image, the return
+        type is a string. If text is a list, a list of strings is returned
+
+        Args:
+            images: image|list
+
+        Returns:
+            list of captions
+        """
+
+        # Convert single element to list
+        values = [images] if not isinstance(images, list) else images
+
+        # Open images if file strings
+        values = [Image.open(image) if isinstance(image, str) else image for image in values]
+
+        # Get and clean captions
+        captions = []
+        for result in self.pipeline(values):
+            text = " ".join([r["generated_text"] for r in result]).strip()
+            captions.append(text)
+
+        # Return single element if single element passed in
+        return captions[0] if not isinstance(images, list) else captions

txtai/pipeline/image/imagehash.py

@@ -0,0 +1,90 @@
+"""
+ImageHash module
+"""
+
+import numpy as np
+
+# Conditional import
+try:
+    from PIL import Image
+    import imagehash
+
+    PIL = True
+except ImportError:
+    PIL = False
+
+from ..base import Pipeline
+
+
+class ImageHash(Pipeline):
+    """
+    Generates perceptual image hashes. These hashes can be used to detect near-duplicate images. This method is not
+    backed by machine learning models and not intended to find conceptually similar images.
+    """
+
+    def __init__(self, algorithm="average", size=8, strings=True):
+        """
+        Creates an ImageHash pipeline.
+
+        Args:
+            algorithm: image hashing algorithm (average, perceptual, difference, wavelet, color)
+            size: hash size
+            strings: outputs hex strings if True (default), otherwise the pipeline returns numpy arrays
+        """
+
+        if not PIL:
+            raise ImportError('ImageHash pipeline is not available - install "pipeline" extra to enable')
+
+        self.algorithm = algorithm
+        self.size = size
+        self.strings = strings
+
+    def __call__(self, images):
+        """
+        Generates perceptual image hashes.
+
+        Args:
+            images: image|list
+
+        Returns:
+            list of hashes
+        """
+
+        # Convert single element to list
+        values = [images] if not isinstance(images, list) else images
+
+        # Open images if file strings
+        values = [Image.open(image) if isinstance(image, str) else image for image in values]
+
+        # Convert images to hashes
+        hashes = [self.ihash(image) for image in values]
+
+        # Return single element if single element passed in
+        return hashes[0] if not isinstance(images, list) else hashes
+
+    def ihash(self, image):
+        """
+        Gets an image hash for image.
+
+        Args:
+            image: PIL image
+
+        Returns:
+            hash as hex string
+        """
+
+        # Apply hash algorithm
+        if self.algorithm == "perceptual":
+            data = imagehash.phash(image, self.size)
+        elif self.algorithm == "difference":
+            data = imagehash.dhash(image, self.size)
+        elif self.algorithm == "wavelet":
+            data = imagehash.whash(image, self.size)
+        elif self.algorithm == "color":
+            data = imagehash.colorhash(image, self.size)
+        else:
+            # Default to average hash
+            data = imagehash.average_hash(image, self.size)
+
+        # Convert to output data type
+        return str(data) if self.strings else data.hash.astype(np.float32).reshape(-1)

txtai/pipeline/image/objects.py

@@ -0,0 +1,80 @@
+"""
+Objects module
+"""
+
+# Conditional import
+try:
+    from PIL import Image
+
+    PIL = True
+except ImportError:
+    PIL = False
+
+from ..hfpipeline import HFPipeline
+
+
+class Objects(HFPipeline):
+    """
+    Applies object detection models to images. Supports both object detection models and image classification models.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, classification=False, threshold=0.9, **kwargs):
+        if not PIL:
+            raise ImportError('Objects pipeline is not available - install "pipeline" extra to enable')
+
+        super().__init__("image-classification" if classification else "object-detection", path, quantize, gpu, model, **kwargs)
+
+        self.classification = classification
+        self.threshold = threshold
+
+    def __call__(self, images, flatten=False, workers=0):
+        """
+        Applies object detection/image classification models to images. Returns a list of (label, score).
+
+        This method supports a single image or a list of images. If the input is an image, the return
+        type is a 1D list of (label, score). If text is a list, a 2D list of (label, score) is
+        returned with a row per image.
+
+        Args:
+            images: image|list
+            flatten: flatten output to a list of objects
+            workers: number of concurrent workers to use for processing data, defaults to None
+
+        Returns:
+            list of (label, score)
+        """
+
+        # Convert single element to list
+        values = [images] if not isinstance(images, list) else images
+
+        # Open images if file strings
+        values = [Image.open(image) if isinstance(image, str) else image for image in values]
+
+        # Run pipeline
+        results = (
+            self.pipeline(values, num_workers=workers)
+            if self.classification
+            else self.pipeline(values, threshold=self.threshold, num_workers=workers)
+        )
+
+        # Build list of (id, score)
+        outputs = []
+        for result in results:
+            # Convert to (label, score) tuples
+            result = [(x["label"], x["score"]) for x in result if x["score"] > self.threshold]
+
+            # Sort by score descending
+            result = sorted(result, key=lambda x: x[1], reverse=True)
+
+            # Deduplicate labels
+            unique = set()
+            elements = []
+            for label, score in result:
+                if label not in unique:
+                    elements.append(label if flatten else (label, score))
+                    unique.add(label)
+
+            outputs.append(elements)
+
+        # Return single element if single element passed in
+        return outputs[0] if not isinstance(images, list) else outputs

txtai/pipeline/llm/__init__.py

@@ -0,0 +1,11 @@
+"""
+LLM imports
+"""
+
+from .factory import GenerationFactory
+from .generation import Generation
+from .huggingface import *
+from .litellm import LiteLLM
+from .llama import LlamaCpp
+from .llm import LLM
+from .rag import RAG

txtai/pipeline/llm/factory.py

@@ -0,0 +1,86 @@
+"""
+Factory module
+"""
+
+from ...util import Resolver
+
+from .huggingface import HFGeneration
+from .litellm import LiteLLM
+from .llama import LlamaCpp
+
+
+class GenerationFactory:
+    """
+    Methods to create generative models.
+    """
+
+    @staticmethod
+    def create(path, method, **kwargs):
+        """
+        Creates a new Generation instance.
+
+        Args:
+            path: model path
+            method: llm framework
+            kwargs: model keyword arguments
+        """
+
+        # Derive method
+        method = GenerationFactory.method(path, method)
+
+        # LiteLLM generation
+        if method == "litellm":
+            return LiteLLM(path, **kwargs)
+
+        # llama.cpp generation
+        if method == "llama.cpp":
+            return LlamaCpp(path, **kwargs)
+
+        # Hugging Face Transformers generation
+        if method == "transformers":
+            return HFGeneration(path, **kwargs)
+
+        # Resolve custom method
+        return GenerationFactory.resolve(path, method, **kwargs)
+
+    @staticmethod
+    def method(path, method):
+        """
+        Get or derives the LLM framework.
+
+        Args:
+            path: model path
+            method: llm framework
+
+        Return:
+            llm framework
+        """
+
+        if not method:
+            if LiteLLM.ismodel(path):
+                method = "litellm"
+            elif LlamaCpp.ismodel(path):
+                method = "llama.cpp"
+            else:
+                method = "transformers"
+
+        return method
+
+    @staticmethod
+    def resolve(path, method, **kwargs):
+        """
+        Attempt to resolve a custom LLM framework.
+
+        Args:
+            path: model path
+            method: llm framework
+            kwargs: model keyword arguments
+
+        Returns:
+            Generation instance
+        """
+
+        try:
+            return Resolver()(method)(path, **kwargs)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve generation framework: '{method}'") from e

txtai/pipeline/llm/generation.py

@@ -0,0 +1,173 @@
+"""
+Generation module
+"""
+
+import re
+
+from ...util import TemplateFormatter
+
+
+class Generation:
+    """
+    Base class for generative models. This class has common logic for building prompts and cleaning model results.
+    """
+
+    def __init__(self, path=None, template=None, **kwargs):
+        """
+        Creates a new Generation instance.
+
+        Args:
+            path: model path
+            template: prompt template
+            kwargs: additional keyword arguments
+        """
+
+        self.path = path
+        self.template = template
+        self.kwargs = kwargs
+
+    def __call__(self, text, maxlength, stream, stop, defaultrole, stripthink, **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
+            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
+        """
+
+        # Format inputs
+        texts = [text] if isinstance(text, str) or isinstance(text[0], dict) else text
+
+        # Apply template, if necessary
+        if self.template:
+            formatter = TemplateFormatter()
+            texts = [formatter.format(self.template, text=x) if isinstance(x, str) else x for x in texts]
+
+        # Apply default role, if necessary
+        if defaultrole == "user":
+            texts = [[{"role": "user", "content": x}] if isinstance(x, str) else x for x in texts]
+
+        # Run pipeline
+        results = self.execute(texts, maxlength, stream, stop, **kwargs)
+
+        # Streaming generation
+        if stream:
+            return results
+
+        # Clean generated text
+        results = [self.clean(texts[x], result, stripthink) for x, result in enumerate(results)]
+
+        # Extract results based on inputs
+        return results[0] if isinstance(text, str) or isinstance(text[0], dict) else results
+
+    def isvision(self):
+        """
+        Returns True if this LLM supports vision operations.
+
+        Returns:
+            True if this is a vision model
+        """
+
+        return False
+
+    def execute(self, texts, maxlength, stream, stop, **kwargs):
+        """
+        Runs a list of prompts through a generative model.
+
+        Args:
+            texts: list of prompts to run
+            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
+        """
+
+        # Streaming generation
+        if stream:
+            return self.stream(texts, maxlength, stream, stop, **kwargs)
+
+        # Full response as content elements
+        return list(self.stream(texts, maxlength, stream, stop, **kwargs))
+
+    def clean(self, prompt, result, stripthink):
+        """
+        Applies a series of rules to clean generated text.
+
+        Args:
+            prompt: original input prompt
+            result: result text
+            stripthink: removes thinking tags if true
+
+        Returns:
+            clean text
+        """
+
+        # Replace input prompt
+        text = result.replace(prompt, "") if isinstance(prompt, str) else result
+
+        # Replace thinking tags, if necessary
+        text = re.sub(r"(?s)<think>.+?</think>", "", text).strip() if stripthink else text
+
+        # Apply text cleaning rules
+        return text.replace("$=", "<=").strip()
+
+    def response(self, result):
+        """
+        Parses response content from the result. This supports both standard and streaming
+        generation.
+
+        For standard generation, the full response is returned. For streaming generation,
+        this method will stream chunks of content.
+
+        Args:
+            result: LLM response
+
+        Returns:
+            response
+        """
+
+        streamed = False
+        for chunk in result:
+            # Expects one of the following parameter paths
+            #  - text
+            #  - message.content
+            #  - delta.content
+            data = chunk["choices"][0]
+            text = data.get("text", data.get("message", data.get("delta")))
+            text = text if isinstance(text, str) else text.get("content")
+
+            # Yield result if there is text AND it's not leading stream whitespace
+            if text is not None and (streamed or text.strip()):
+                yield (text.lstrip() if not streamed else text)
+                streamed = True
+
+    def stream(self, texts, maxlength, stream, stop, **kwargs):
+        """
+        Streams LLM responses.
+
+        Args:
+            texts: list of prompts to run
+            maxlength: maximum sequence length
+            stream: stream response if True, defaults to False
+            stop: list of stop strings
+            kwargs: additional generation keyword arguments
+
+        Returns:
+            responses
+        """
+
+        raise NotImplementedError