mseep-txtai 9.1.1__py3-none-any.whl

txtai/pipeline/text/entity.py

@@ -0,0 +1,140 @@
+"""
+Entity module
+"""
+
+# Conditional import
+try:
+    from gliner import GLiNER
+
+    GLINER = True
+except ImportError:
+    GLINER = False
+
+from huggingface_hub.errors import HFValidationError
+from transformers.utils import cached_file
+
+from ...models import Models
+from ..hfpipeline import HFPipeline
+
+
+class Entity(HFPipeline):
+    """
+    Applies a token classifier to text and extracts entity/label combinations.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs):
+        # Create a new entity pipeline
+        self.gliner = self.isgliner(path)
+        if self.gliner:
+            if not GLINER:
+                raise ImportError('GLiNER is not available - install "pipeline" extra to enable')
+
+            # GLiNER entity pipeline
+            self.pipeline = GLiNER.from_pretrained(path)
+            self.pipeline = self.pipeline.to(Models.device(Models.deviceid(gpu)))
+        else:
+            # Standard entity pipeline
+            super().__init__("token-classification", path, quantize, gpu, model, **kwargs)
+
+    def __call__(self, text, labels=None, aggregate="simple", flatten=None, join=False, workers=0):
+        """
+        Applies a token classifier to text and extracts entity/label combinations.
+
+        Args:
+            text: text|list
+            labels: list of entity type labels to accept, defaults to None which accepts all
+            aggregate: method to combine multi token entities - options are "simple" (default), "first", "average" or "max"
+            flatten: flatten output to a list of labels if present. Accepts a boolean or float value to only keep scores greater than that number.
+            join: joins flattened output into a string if True, ignored if flatten not set
+            workers: number of concurrent workers to use for processing data, defaults to None
+
+        Returns:
+            list of (entity, entity type, score) or list of entities depending on flatten parameter
+        """
+
+        # Run token classification pipeline
+        results = self.execute(text, labels, aggregate, workers)
+
+        # Convert results to a list if necessary
+        if isinstance(text, str):
+            results = [results]
+
+        # Score threshold when flatten is set
+        threshold = 0.0 if isinstance(flatten, bool) else flatten
+
+        # Extract entities if flatten set, otherwise extract (entity, entity type, score) tuples
+        outputs = []
+        for result in results:
+            if flatten:
+                output = [r["word"] for r in result if self.accept(r["entity_group"], labels) and r["score"] >= threshold]
+                outputs.append(" ".join(output) if join else output)
+            else:
+                outputs.append([(r["word"], r["entity_group"], float(r["score"])) for r in result if self.accept(r["entity_group"], labels)])
+
+        return outputs[0] if isinstance(text, str) else outputs
+
+    def isgliner(self, path):
+        """
+        Tests if path is a GLiNER model.
+
+        Args:
+            path: model path
+
+        Returns:
+            True if this is a GLiNER model, False otherwise
+        """
+
+        try:
+            # Test if this model has a gliner_config.json file
+            return cached_file(path_or_repo_id=path, filename="gliner_config.json") is not None
+
+        # Ignore this error - invalid repo or directory
+        except (HFValidationError, OSError):
+            pass
+
+        return False
+
+    def execute(self, text, labels, aggregate, workers):
+        """
+        Runs the entity extraction pipeline.
+
+        Args:
+            text: text|list
+            labels: list of entity type labels to accept, defaults to None which accepts all
+            aggregate: method to combine multi token entities - options are "simple" (default), "first", "average" or "max"
+            workers: number of concurrent workers to use for processing data, defaults to None
+
+        Returns:
+            list of entities and labels
+        """
+
+        if self.gliner:
+            # Extract entities with GLiNER. Use default CoNLL-2003 labels when not otherwise provided.
+            results = self.pipeline.batch_predict_entities(
+                text if isinstance(text, list) else [text], labels if labels else ["person", "organization", "location"]
+            )
+
+            # Map results to same format as Transformers token classifier
+            entities = []
+            for result in results:
+                entities.append([{"word": x["text"], "entity_group": x["label"], "score": x["score"]} for x in result])
+
+            # Return extracted entities
+            return entities if isinstance(text, list) else entities[0]
+
+        # Standard Transformers token classification pipeline
+        return self.pipeline(text, aggregation_strategy=aggregate, num_workers=workers)
+
+    def accept(self, etype, labels):
+        """
+        Determines if entity type is in valid entity type.
+
+        Args:
+            etype: entity type
+            labels: list of entities to accept
+
+        Returns:
+            if etype is accepted
+        """
+
+        return not labels or etype in labels

txtai/pipeline/text/labels.py

@@ -0,0 +1,137 @@
+"""
+Labels module
+"""
+
+from ..hfpipeline import HFPipeline
+
+
+class Labels(HFPipeline):
+    """
+    Applies a text classifier to text. Supports zero shot and standard text classification models
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, dynamic=True, **kwargs):
+        super().__init__("zero-shot-classification" if dynamic else "text-classification", path, quantize, gpu, model, **kwargs)
+
+        # Set if labels are dynamic (zero shot) or fixed (standard text classification)
+        self.dynamic = dynamic
+
+    def __call__(self, text, labels=None, multilabel=False, flatten=None, workers=0, **kwargs):
+        """
+        Applies a text classifier to text. Returns a list of (id, score) sorted by highest score,
+        where id is the index in labels. For zero shot classification, a list of labels is required.
+        For text classification models, a list of labels is optional, otherwise all trained labels are returned.
+
+        This method supports text as a string or a list. If the input is a string, the return
+        type is a 1D list of (id, score). If text is a list, a 2D list of (id, score) is
+        returned with a row per string.
+
+        Args:
+            text: text|list
+            labels: list of labels
+            multilabel: labels are independent if True, scores are normalized to sum to 1 per text item if False, raw scores returned if None
+            flatten: flatten output to a list of labels if present. Accepts a boolean or float value to only keep scores greater than that number.
+            workers: number of concurrent workers to use for processing data, defaults to None
+            kwargs: additional keyword args
+
+        Returns:
+            list of (id, score) or list of labels depending on flatten parameter
+        """
+
+        if self.dynamic:
+            # Run zero shot classification pipeline
+            results = self.pipeline(text, labels, multi_label=multilabel, truncation=True, num_workers=workers)
+        else:
+            # Set classification function based on inputs
+            function = "none" if multilabel is None else "sigmoid" if multilabel or len(self.labels()) == 1 else "softmax"
+
+            # Run text classification pipeline
+            results = self.pipeline(text, top_k=None, function_to_apply=function, num_workers=workers, **kwargs)
+
+        # Convert results to a list if necessary
+        if isinstance(text, str):
+            results = [results]
+
+        # Build list of outputs and return
+        outputs = self.outputs(results, labels, flatten)
+        return outputs[0] if isinstance(text, str) else outputs
+
+    def labels(self):
+        """
+        Returns a list of all text classification model labels sorted in index order.
+
+        Returns:
+            list of labels
+        """
+
+        return list(self.pipeline.model.config.id2label.values())
+
+    def outputs(self, results, labels, flatten):
+        """
+        Processes pipeline results and builds outputs.
+
+        Args:
+            results: pipeline results
+            labels: list of labels
+            flatten: flatten output to a list of labels if present. Accepts a boolean or float value to only keep scores greater than that number.
+
+        Returns:
+            list of outputs
+        """
+
+        outputs = []
+        threshold = 0.0 if isinstance(flatten, bool) else flatten
+
+        for result in results:
+            if self.dynamic:
+                if flatten:
+                    result = [label for x, label in enumerate(result["labels"]) if result["scores"][x] >= threshold]
+                    outputs.append(result[:1] if isinstance(flatten, bool) else result)
+                else:
+                    outputs.append([(labels.index(label), result["scores"][x]) for x, label in enumerate(result["labels"])])
+            else:
+                if flatten:
+                    result = [x["label"] for x in result if x["score"] >= threshold and (not labels or x["label"] in labels)]
+                    outputs.append(result[:1] if isinstance(flatten, bool) else result)
+                else:
+                    # Filter results using labels, if provided
+                    outputs.append(self.limit(result, labels))
+
+        return outputs
+
+    def limit(self, result, labels):
+        """
+        Filter result using labels. If labels is None, original result is returned.
+
+        Args:
+            result: results array sorted by score descending
+            labels: list of labels or None
+
+        Returns:
+            filtered results
+        """
+
+        # Get config
+        config = self.pipeline.model.config
+
+        # Resolve label ids for labels
+        result = [(config.label2id.get(x["label"], 0), x["score"]) for x in result]
+
+        if labels:
+            matches = []
+            for label in labels:
+                # Lookup label keys from model config
+                if label.isdigit():
+                    label = int(label)
+                    keys = list(config.id2label.keys())
+                else:
+                    label = label.lower()
+                    keys = [x.lower() for x in config.label2id.keys()]
+
+                # Find and add label match
+                if label in keys:
+                    matches.append(keys.index(label))
+
+            return [(label, score) for label, score in result if label in matches]
+
+        return result

txtai/pipeline/text/lateencoder.py

@@ -0,0 +1,103 @@
+"""
+Late encoder module
+"""
+
+import numpy as np
+import torch
+
+from ...models import Models, PoolingFactory
+from ..base import Pipeline
+
+
+class LateEncoder(Pipeline):
+    """
+    Computes similarity between query and list of text using a late interaction model.
+    """
+
+    def __init__(self, path=None, **kwargs):
+        # Get device
+        self.device = Models.device(Models.deviceid(kwargs.get("gpu", True)))
+
+        # Load model
+        self.model = PoolingFactory.create(
+            {
+                "method": kwargs.get("method"),
+                "path": path if path else "colbert-ir/colbertv2.0",
+                "device": self.device,
+                "tokenizer": kwargs.get("tokenizer"),
+                "maxlength": kwargs.get("maxlength"),
+                "modelargs": {**kwargs.get("vectors", {}), **{"muvera": None}},
+            }
+        )
+
+    def __call__(self, query, texts, limit=None):
+        """
+        Computes the similarity between query and list of text. Returns a list of
+        (id, score) sorted by highest score, where id is the index in texts.
+
+        This method supports query as a string or a list. If the input is a string,
+        the return type is a 1D list of (id, score). If text is a list, a 2D list
+        of (id, score) is returned with a row per string.
+
+        Args:
+            query: query text|list
+            texts: list of text
+            limit: maximum comparisons to return, defaults to all
+
+        Returns:
+            list of (id, score)
+        """
+
+        queries = [query] if isinstance(query, str) else query
+
+        # Encode text to vectors
+        queries = self.encode(queries, "query")
+        data = self.encode(texts, "data") if isinstance(texts[0], str) else texts
+
+        # Compute maximum similarity score
+        scores = []
+        for q in queries:
+            scores.extend(self.score(q.unsqueeze(0), data, limit))
+
+        return scores[0] if isinstance(query, str) else scores
+
+    def encode(self, data, category):
+        """
+        Encodes a batch of data using the underlying model.
+
+        Args:
+            data: input data
+            category: encoding category
+
+        Returns:
+            encoded data
+        """
+
+        return torch.from_numpy(self.model.encode(data, category=category)).to(self.device)
+
+    def score(self, queries, data, limit):
+        """
+        Computes the maximum similarity score between query vectors and data vectors.
+
+        Args:
+            queries: query vectors
+            data: data vectors
+            limit: query limit
+
+        Returns:
+            list of (id, score)
+        """
+
+        # Compute bulk dot product using einstein notation
+        scores = torch.einsum("ash,bth->abst", queries, data).max(axis=-1).values.mean(axis=-1)
+        scores = scores.cpu().numpy()
+
+        # Get top n matching indices and scores
+        indices = np.argpartition(-scores, limit if limit and limit < scores.shape[0] else scores.shape[0] - 1)[:, :limit]
+        scores = np.take_along_axis(scores, indices, axis=1)
+
+        results = []
+        for x, index in enumerate(indices):
+            results.append(list(zip(index.tolist(), scores[x].tolist())))
+
+        return results

txtai/pipeline/text/questions.py

@@ -0,0 +1,48 @@
+"""
+Questions module
+"""
+
+from ..hfpipeline import HFPipeline
+
+
+class Questions(HFPipeline):
+    """
+    Runs extractive QA for a series of questions and contexts.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs):
+        super().__init__("question-answering", path, quantize, gpu, model, **kwargs)
+
+    def __call__(self, questions, contexts, workers=0):
+        """
+        Runs a extractive question-answering model against each question-context pair, finding the best answers.
+
+        Args:
+            questions: list of questions
+            contexts: list of contexts to pull answers from
+            workers: number of concurrent workers to use for processing data, defaults to None
+
+        Returns:
+            list of answers
+        """
+
+        answers = []
+
+        for x, question in enumerate(questions):
+            if question and contexts[x]:
+                # Run the QA pipeline
+                result = self.pipeline(question=question, context=contexts[x], num_workers=workers)
+
+                # Get answer and score
+                answer, score = result["answer"], result["score"]
+
+                # Require score to be at least 0.05
+                if score < 0.05:
+                    answer = None
+
+                # Add answer
+                answers.append(answer)
+            else:
+                answers.append(None)
+
+        return answers

txtai/pipeline/text/reranker.py

@@ -0,0 +1,57 @@
+"""
+Reranker module
+"""
+
+from ..base import Pipeline
+
+
+class Reranker(Pipeline):
+    """
+    Runs embeddings queries and re-ranks them using a similarity pipeline. Note that content must be enabled with the
+    embeddings instance for this to work properly.
+    """
+
+    def __init__(self, embeddings, similarity):
+        """
+        Creates a Reranker pipeline.
+
+        Args:
+            embeddings: embeddings instance (content must be enabled)
+            similarity: similarity instance
+        """
+
+        self.embeddings, self.similarity = embeddings, similarity
+
+    # pylint: disable=W0222
+    def __call__(self, query, limit=3, factor=10, **kwargs):
+        """
+        Runs an embeddings search and re-ranks the results using a Similarity pipeline.
+
+        Args:
+            query: query text|list
+            limit: maximum results
+            factor: factor to multiply limit by for the initial embeddings search
+            kwargs: additional arguments to pass to embeddings search
+
+        Returns:
+            list of query results rescored using a Similarity pipeline
+        """
+
+        queries = [query] if not isinstance(query, list) else query
+
+        # Run searches
+        results = self.embeddings.batchsearch(queries, limit * factor, **kwargs)
+
+        # Re-rank using similarity pipeline
+        ranked = []
+        for x, result in enumerate(results):
+            texts = [row["text"] for row in result]
+
+            # Score results and merge
+            for uid, score in self.similarity(queries[x], texts):
+                result[uid]["score"] = score
+
+            # Sort and take top n sorted results
+            ranked.append(sorted(result, key=lambda row: row["score"], reverse=True)[:limit])
+
+        return ranked[0] if isinstance(query, str) else ranked

txtai/pipeline/text/similarity.py

@@ -0,0 +1,83 @@
+"""
+Similarity module
+"""
+
+import numpy as np
+
+from .crossencoder import CrossEncoder
+from .labels import Labels
+from .lateencoder import LateEncoder
+
+
+class Similarity(Labels):
+    """
+    Computes similarity between query and list of text using a transformers model.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, dynamic=True, crossencode=False, lateencode=False, **kwargs):
+        self.crossencoder, self.lateencoder = None, None
+
+        if lateencode:
+            # Load a late interaction encoder if lateencode set to True
+            self.lateencoder = LateEncoder(path=path, gpu=gpu, **kwargs)
+        else:
+            # Use zero-shot classification if dynamic is True and crossencode is False, otherwise use standard text classification
+            super().__init__(path, quantize, gpu, model, False if crossencode else dynamic, **kwargs)
+
+            # Load as a cross-encoder if crossencode set to True
+            self.crossencoder = CrossEncoder(model=self.pipeline) if crossencode else None
+
+    # pylint: disable=W0222
+    def __call__(self, query, texts, multilabel=True, **kwargs):
+        """
+        Computes the similarity between query and list of text. Returns a list of
+        (id, score) sorted by highest score, where id is the index in texts.
+
+        This method supports query as a string or a list. If the input is a string,
+        the return type is a 1D list of (id, score). If text is a list, a 2D list
+        of (id, score) is returned with a row per string.
+
+        Args:
+            query: query text|list
+            texts: list of text
+            multilabel: labels are independent if True, scores are normalized to sum to 1 per text item if False, raw scores returned if None
+            kwargs: additional keyword args
+
+        Returns:
+            list of (id, score)
+        """
+
+        if self.crossencoder:
+            # pylint: disable=E1102
+            return self.crossencoder(query, texts, multilabel)
+
+        if self.lateencoder:
+            return self.lateencoder(query, texts)
+
+        # Call Labels pipeline for texts using input query as the candidate label
+        scores = super().__call__(texts, [query] if isinstance(query, str) else query, multilabel, **kwargs)
+
+        # Sort on query index id
+        scores = [[score for _, score in sorted(row)] for row in scores]
+
+        # Transpose axes to get a list of text scores for each query
+        scores = np.array(scores).T.tolist()
+
+        # Build list of (id, score) per query sorted by highest score
+        scores = [sorted(enumerate(row), key=lambda x: x[1], reverse=True) for row in scores]
+
+        return scores[0] if isinstance(query, str) else scores
+
+    def encode(self, data, category):
+        """
+        Encodes a batch of data using the underlying model.
+
+        Args:
+            data: input data
+            category: encoding category
+
+        Returns:
+            encoded data
+        """
+
+        return self.lateencoder.encode(data, category) if self.lateencoder else data

txtai/pipeline/text/summary.py

@@ -0,0 +1,98 @@
+"""
+Summary module
+"""
+
+import re
+
+from ..hfpipeline import HFPipeline
+
+
+class Summary(HFPipeline):
+    """
+    Summarizes text.
+    """
+
+    def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs):
+        super().__init__("summarization", path, quantize, gpu, model, **kwargs)
+
+    def __call__(self, text, minlength=None, maxlength=None, workers=0):
+        """
+        Runs a summarization model against a block of text.
+
+        This method supports text as a string or a list. If the input is a string, the return
+        type is text. If text is a list, a list of text is returned with a row per block of text.
+
+        Args:
+            text: text|list
+            minlength: minimum length for summary
+            maxlength: maximum length for summary
+            workers: number of concurrent workers to use for processing data, defaults to None
+
+        Returns:
+            summary text
+        """
+
+        # Validate text length greater than max length
+        check = maxlength if maxlength else self.maxlength()
+
+        # Skip text shorter than max length
+        texts = text if isinstance(text, list) else [text]
+        params = [(x, text if len(text) >= check else None) for x, text in enumerate(texts)]
+
+        # Build keyword arguments
+        kwargs = self.args(minlength, maxlength)
+
+        inputs = [text for _, text in params if text]
+        if inputs:
+            # Run summarization pipeline
+            results = self.pipeline(inputs, num_workers=workers, **kwargs)
+
+            # Pull out summary text
+            results = iter([self.clean(x["summary_text"]) for x in results])
+            results = [next(results) if text else texts[x] for x, text in params]
+        else:
+            # Return original
+            results = texts
+
+        return results[0] if isinstance(text, str) else results
+
+    def clean(self, text):
+        """
+        Applies a series of rules to clean extracted text.
+
+        Args:
+            text: input text
+
+        Returns:
+            clean text
+        """
+
+        text = re.sub(r"\s*\.\s*", ". ", text)
+        text = text.strip()
+
+        return text
+
+    def args(self, minlength, maxlength):
+        """
+        Builds keyword arguments.
+
+        Args:
+            minlength: minimum length for summary
+            maxlength: maximum length for summary
+
+        Returns:
+            keyword arguments
+        """
+
+        kwargs = {"truncation": True}
+        if minlength:
+            kwargs["min_length"] = minlength
+        if maxlength:
+            kwargs["max_length"] = maxlength
+            kwargs["max_new_tokens"] = None
+
+            # Default minlength if not provided or it's bigger than maxlength
+            if "min_length" not in kwargs or kwargs["min_length"] > kwargs["max_length"]:
+                kwargs["min_length"] = kwargs["max_length"]
+
+        return kwargs