mseep-txtai 9.1.1__py3-none-any.whl

txtai/graph/topics.py

@@ -0,0 +1,166 @@
+"""
+Topics module
+"""
+
+from ..pipeline import Tokenizer
+from ..scoring import ScoringFactory
+
+
+class Topics:
+    """
+    Topic modeling using community detection.
+    """
+
+    def __init__(self, config):
+        """
+        Creates a new Topics instance.
+
+        Args:
+            config: topic configuration
+        """
+
+        self.config = config if config else {}
+        self.tokenizer = Tokenizer(stopwords=True)
+
+        # Additional stopwords to ignore when building topic names
+        self.stopwords = set()
+        if "stopwords" in self.config:
+            self.stopwords.update(self.config["stopwords"])
+
+    def __call__(self, graph):
+        """
+        Runs topic modeling for input graph.
+
+        Args:
+            graph: Graph instance
+
+        Returns:
+            dictionary of {topic name: [ids]}
+        """
+
+        # Detect communities
+        communities = graph.communities(self.config)
+
+        # Sort by community size, largest to smallest
+        communities = sorted(communities, key=len, reverse=True)
+
+        # Calculate centrality of graph
+        centrality = graph.centrality()
+
+        # Score communities and generate topn terms
+        topics = [self.score(graph, x, community, centrality) for x, community in enumerate(communities)]
+
+        # Merge duplicate topics and return
+        return self.merge(topics)
+
+    def score(self, graph, index, community, centrality):
+        """
+        Scores a community of nodes and generates the topn terms in the community.
+
+        Args:
+            graph: Graph instance
+            index: community index
+            community: community of nodes
+            centrality: node centrality scores
+
+        Returns:
+            (topn topic terms, topic ids sorted by score descending)
+        """
+
+        # Tokenize input and build scoring index
+        scoring = ScoringFactory.create({"method": self.config.get("labels", "bm25"), "terms": True})
+        scoring.index(((node, self.tokenize(graph, node), None) for node in community))
+
+        # Check if scoring index has data
+        if scoring.idf:
+            # Sort by most commonly occurring terms (i.e. lowest score)
+            idf = sorted(scoring.idf, key=scoring.idf.get)
+
+            # Term count for generating topic labels
+            topn = self.config.get("terms", 4)
+
+            # Get topn terms
+            terms = self.topn(idf, topn)
+
+            # Sort community by score descending
+            community = [uid for uid, _ in scoring.search(terms, len(community))]
+        else:
+            # No text found for topic, generate topic name
+            terms = ["topic", str(index)]
+
+            # Sort community by centrality scores
+            community = sorted(community, key=lambda x: centrality[x], reverse=True)
+
+        return (terms, community)
+
+    def tokenize(self, graph, node):
+        """
+        Tokenizes node text.
+
+        Args:
+            graph: Graph instance
+            node: node id
+
+        Returns:
+            list of node tokens
+        """
+
+        text = graph.attribute(node, "text")
+        return self.tokenizer(text) if text else []
+
+    def topn(self, terms, n):
+        """
+        Gets topn terms.
+
+        Args:
+            terms: list of terms
+            n: topn
+
+        Returns:
+            topn terms
+        """
+
+        topn = []
+
+        for term in terms:
+            # Add terms that pass tokenization rules
+            if self.tokenizer(term) and term not in self.stopwords:
+                topn.append(term)
+
+            # Break once topn terms collected
+            if len(topn) == n:
+                break
+
+        return topn
+
+    def merge(self, topics):
+        """
+        Merges duplicate topics
+
+        Args:
+            topics: list of (topn terms, topic ids)
+
+        Returns:
+            dictionary of {topic name:[ids]}
+        """
+
+        merge, termslist = {}, {}
+
+        for terms, uids in topics:
+            # Use topic terms as key
+            key = frozenset(terms)
+
+            # Add key to merged topics, if necessary
+            if key not in merge:
+                merge[key], termslist[key] = [], terms
+
+            # Merge communities
+            merge[key].extend(uids)
+
+        # Sort communities largest to smallest since the order could have changed with merges
+        results = {}
+        for k, v in sorted(merge.items(), key=lambda x: len(x[1]), reverse=True):
+            # Create composite string key using topic terms and store ids
+            results["_".join(termslist[k])] = v
+
+        return results

txtai/models/__init__.py

@@ -0,0 +1,9 @@
+"""
+Models imports
+"""
+
+from .models import Models
+from .onnx import OnnxModel
+from .pooling import *
+from .registry import Registry
+from .tokendetection import TokenDetection

txtai/models/models.py

@@ -0,0 +1,268 @@
+"""
+Models module
+"""
+
+import os
+
+import torch
+
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForQuestionAnswering,
+    AutoModelForSeq2SeqLM,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+)
+from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
+
+from .onnx import OnnxModel
+
+
+class Models:
+    """
+    Utility methods for working with machine learning models
+    """
+
+    @staticmethod
+    def checklength(config, tokenizer):
+        """
+        Checks the length for a Hugging Face Transformers tokenizer using a Hugging Face Transformers config. Copies the
+        max_position_embeddings parameter if the tokenizer has no max_length set. This helps with backwards compatibility
+        with older tokenizers.
+
+        Args:
+            config: transformers config
+            tokenizer: transformers tokenizer
+        """
+
+        # Unpack nested config, handles passing model directly
+        if hasattr(config, "config"):
+            config = config.config
+
+        if (
+            hasattr(config, "max_position_embeddings")
+            and tokenizer
+            and hasattr(tokenizer, "model_max_length")
+            and tokenizer.model_max_length == int(1e30)
+        ):
+            tokenizer.model_max_length = config.max_position_embeddings
+
+    @staticmethod
+    def maxlength(config, tokenizer):
+        """
+        Gets the best max length to use for generate calls. This method will return config.max_length if it's set. Otherwise, it will return
+        tokenizer.model_max_length.
+
+        Args:
+            config: transformers config
+            tokenizer: transformers tokenizer
+        """
+
+        # Unpack nested config, handles passing model directly
+        if hasattr(config, "config"):
+            config = config.config
+
+        # Get non-defaulted fields
+        keys = config.to_diff_dict()
+
+        # Use config.max_length if not set to default value, else use tokenizer.model_max_length if available
+        return config.max_length if "max_length" in keys or not hasattr(tokenizer, "model_max_length") else tokenizer.model_max_length
+
+    @staticmethod
+    def deviceid(gpu):
+        """
+        Translates input gpu argument into a device id.
+
+        Args:
+            gpu: True/False if GPU should be enabled, also supports a device id/string/instance
+
+        Returns:
+            device id
+        """
+
+        # Return if this is already a torch device
+        # pylint: disable=E1101
+        if isinstance(gpu, torch.device):
+            return gpu
+
+        # Always return -1 if gpu is None or an accelerator device is unavailable
+        if gpu is None or not Models.hasaccelerator():
+            return -1
+
+        # Default to device 0 if gpu is True and not otherwise specified
+        if isinstance(gpu, bool):
+            return 0 if gpu else -1
+
+        # Return gpu as device id if gpu flag is an int
+        return int(gpu)
+
+    @staticmethod
+    def device(deviceid):
+        """
+        Gets a tensor device.
+
+        Args:
+            deviceid: device id
+
+        Returns:
+            tensor device
+        """
+
+        # Torch device
+        # pylint: disable=E1101
+        return deviceid if isinstance(deviceid, torch.device) else torch.device(Models.reference(deviceid))
+
+    @staticmethod
+    def reference(deviceid):
+        """
+        Gets a tensor device reference.
+
+        Args:
+            deviceid: device id
+
+        Returns:
+            device reference
+        """
+
+        return (
+            deviceid
+            if isinstance(deviceid, str)
+            else (
+                "cpu"
+                if deviceid < 0
+                else f"cuda:{deviceid}" if torch.cuda.is_available() else "mps" if Models.hasmpsdevice() else Models.finddevice()
+            )
+        )
+
+    @staticmethod
+    def acceleratorcount():
+        """
+        Gets the number of accelerator devices available.
+
+        Returns:
+            number of accelerators available
+        """
+
+        return max(torch.cuda.device_count(), int(Models.hasaccelerator()))
+
+    @staticmethod
+    def hasaccelerator():
+        """
+        Checks if there is an accelerator device available.
+
+        Returns:
+            True if an accelerator device is available, False otherwise
+        """
+
+        return torch.cuda.is_available() or Models.hasmpsdevice() or bool(Models.finddevice())
+
+    @staticmethod
+    def hasmpsdevice():
+        """
+        Checks if there is a MPS device available.
+
+        Returns:
+            True if a MPS device is available, False otherwise
+        """
+
+        return os.environ.get("PYTORCH_MPS_DISABLE") != "1" and torch.backends.mps.is_available()
+
+    @staticmethod
+    def finddevice():
+        """
+        Attempts to find an alternative accelerator device.
+
+        Returns:
+            name of first alternative accelerator available or None if not found
+        """
+
+        return next((device for device in ["xpu"] if hasattr(torch, device) and getattr(torch, device).is_available()), None)
+
+    @staticmethod
+    def load(path, config=None, task="default", modelargs=None):
+        """
+        Loads a machine learning model. Handles multiple model frameworks (ONNX, Transformers).
+
+        Args:
+            path: path to model
+            config: path to model configuration
+            task: task name used to lookup model type
+
+        Returns:
+            machine learning model
+        """
+
+        # Detect ONNX models
+        if isinstance(path, bytes) or (isinstance(path, str) and os.path.isfile(path)):
+            return OnnxModel(path, config)
+
+        # Return path, if path isn't a string
+        if not isinstance(path, str):
+            return path
+
+        # Transformer models
+        models = {
+            "default": AutoModel.from_pretrained,
+            "question-answering": AutoModelForQuestionAnswering.from_pretrained,
+            "summarization": AutoModelForSeq2SeqLM.from_pretrained,
+            "text-classification": AutoModelForSequenceClassification.from_pretrained,
+            "zero-shot-classification": AutoModelForSequenceClassification.from_pretrained,
+        }
+
+        # Pass modelargs as keyword arguments
+        modelargs = modelargs if modelargs else {}
+
+        # Load model for supported tasks. Return path for unsupported tasks.
+        return models[task](path, **modelargs) if task in models else path
+
+    @staticmethod
+    def tokenizer(path, **kwargs):
+        """
+        Loads a tokenizer from path.
+
+        Args:
+            path: path to tokenizer
+            kwargs: optional additional keyword arguments
+
+        Returns:
+            tokenizer
+        """
+
+        return AutoTokenizer.from_pretrained(path, **kwargs) if isinstance(path, str) else path
+
+    @staticmethod
+    def task(path, **kwargs):
+        """
+        Attempts to detect the model task from path.
+
+        Args:
+            path: path to model
+            kwargs: optional additional keyword arguments
+
+        Returns:
+            inferred model task
+        """
+
+        # Get model configuration
+        config = None
+        if isinstance(path, (list, tuple)) and hasattr(path[0], "config"):
+            config = path[0].config
+        elif isinstance(path, str):
+            config = AutoConfig.from_pretrained(path, **kwargs)
+
+        # Attempt to resolve task using configuration
+        task = None
+        if config:
+            architecture = config.architectures[0] if config.architectures else None
+            if architecture:
+                if architecture in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values():
+                    task = "vision"
+                elif any(x for x in ["LMHead", "CausalLM"] if x in architecture):
+                    task = "language-generation"
+                elif "QuestionAnswering" in architecture:
+                    task = "question-answering"
+                elif "ConditionalGeneration" in architecture:
+                    task = "sequence-sequence"
+
+        return task

txtai/models/onnx.py

@@ -0,0 +1,133 @@
+"""
+ONNX module
+"""
+
+# Conditional import
+try:
+    import onnxruntime as ort
+
+    ONNX_RUNTIME = True
+except ImportError:
+    ONNX_RUNTIME = False
+
+import numpy as np
+import torch
+
+from transformers import AutoConfig
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+
+from .registry import Registry
+
+
+# pylint: disable=W0223
+class OnnxModel(PreTrainedModel):
+    """
+    Provides a Transformers/PyTorch compatible interface for ONNX models. Handles casting inputs
+    and outputs with minimal to no copying of data.
+    """
+
+    def __init__(self, model, config=None):
+        """
+        Creates a new OnnxModel.
+
+        Args:
+            model: path to model or InferenceSession
+            config: path to model configuration
+        """
+
+        if not ONNX_RUNTIME:
+            raise ImportError('onnxruntime is not available - install "model" extra to enable')
+
+        super().__init__(AutoConfig.from_pretrained(config) if config else OnnxConfig())
+
+        # Create ONNX session
+        self.model = ort.InferenceSession(model, ort.SessionOptions(), self.providers())
+
+        # Add references for this class to supported AutoModel classes
+        Registry.register(self)
+
+    @property
+    def device(self):
+        """
+        Returns model device id.
+
+        Returns:
+            model device id
+        """
+
+        return -1
+
+    def providers(self):
+        """
+        Returns a list of available and usable providers.
+
+        Returns:
+            list of available and usable providers
+        """
+
+        # Create list of providers, prefer CUDA provider if available
+        # CUDA provider only available if GPU is available and onnxruntime-gpu installed
+        if torch.cuda.is_available() and "CUDAExecutionProvider" in ort.get_available_providers():
+            return ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        # Default when CUDA provider isn't available
+        return ["CPUExecutionProvider"]
+
+    def forward(self, **inputs):
+        """
+        Runs inputs through an ONNX model and returns outputs. This method handles casting inputs
+        and outputs between torch tensors and numpy arrays as shared memory (no copy).
+
+        Args:
+            inputs: model inputs
+
+        Returns:
+            model outputs
+        """
+
+        inputs = self.parse(inputs)
+
+        # Run inputs through ONNX model
+        results = self.model.run(None, inputs)
+
+        # pylint: disable=E1101
+        # Detect if logits is an output and return classifier output in that case
+        if any(x.name for x in self.model.get_outputs() if x.name == "logits"):
+            return SequenceClassifierOutput(logits=torch.from_numpy(np.array(results[0])))
+
+        return torch.from_numpy(np.array(results))
+
+    def parse(self, inputs):
+        """
+        Parse model inputs and handle converting to ONNX compatible inputs.
+
+        Args:
+            inputs: model inputs
+
+        Returns:
+            ONNX compatible model inputs
+        """
+
+        features = {}
+
+        # Select features from inputs
+        for key in ["input_ids", "attention_mask", "token_type_ids"]:
+            if key in inputs:
+                value = inputs[key]
+
+                # Cast torch tensors to numpy
+                if hasattr(value, "cpu"):
+                    value = value.cpu().numpy()
+
+                # Cast to numpy array if not already one
+                features[key] = np.asarray(value)
+
+        return features
+
+
+class OnnxConfig(PretrainedConfig):
+    """
+    Configuration for ONNX models.
+    """

txtai/models/pooling/__init__.py

@@ -0,0 +1,9 @@
+"""
+Pooling imports
+"""
+
+from .base import Pooling
+from .cls import ClsPooling
+from .factory import PoolingFactory
+from .late import LatePooling
+from .mean import MeanPooling