mseep-txtai 9.1.1__py3-none-any.whl

txtai/console/__main__.py

@@ -0,0 +1,22 @@
+"""
+Main module.
+"""
+
+import sys
+
+from .base import Console
+
+
+def main(path=None):
+    """
+    Console execution loop.
+
+    Args:
+        path: model path
+    """
+
+    Console(path).cmdloop()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1] if len(sys.argv) > 1 else None)

txtai/console/base.py

@@ -0,0 +1,264 @@
+"""
+Console module
+"""
+
+import os
+import shlex
+
+from cmd import Cmd
+
+# Conditional import
+try:
+    from rich import box
+    from rich.console import Console as RichConsole
+    from rich.table import Table
+
+    RICH = True
+except ImportError:
+    RICH = False
+
+from txtai.app import Application
+from txtai.embeddings import Embeddings
+
+
+class Console(Cmd):
+    """
+    txtai console.
+    """
+
+    def __init__(self, path=None):
+        """
+        Creates a new command line console.
+
+        Args:
+            path: path to initial configuration, if any
+        """
+
+        super().__init__()
+
+        if not RICH:
+            raise ImportError('Console is not available - install "console" extra to enable')
+
+        self.prompt = ">>> "
+
+        # Rich console
+        self.console = RichConsole()
+
+        # App parameters
+        self.app = None
+        self.path = path
+
+        # Parameters
+        self.vhighlight = None
+        self.vlimit = None
+
+    def preloop(self):
+        """
+        Loads initial configuration.
+        """
+
+        self.console.print("txtai console", style="#03a9f4")
+
+        # Load default path
+        if self.path:
+            self.load(self.path)
+
+    def default(self, line):
+        """
+        Default event loop.
+
+        Args:
+            line: command line
+        """
+
+        # pylint: disable=W0703
+        try:
+            command = line.lower()
+            if command.startswith(".config"):
+                self.config()
+            elif command.startswith(".highlight"):
+                self.highlight(command)
+            elif command.startswith(".limit"):
+                self.limit(command)
+            elif command.startswith(".load"):
+                command = self.split(line)
+                self.path = command[1]
+                self.load(self.path)
+            elif command.startswith(".workflow"):
+                self.workflow(line)
+            else:
+                # Search is default action
+                self.search(line)
+        except Exception:
+            self.console.print_exception()
+
+    def config(self):
+        """
+        Processes .config command.
+        """
+
+        self.console.print(self.app.config)
+
+    def highlight(self, command):
+        """
+        Processes .highlight command.
+
+        Args:
+            command: command line
+        """
+
+        _, action = self.split(command, "#ffff00")
+        self.vhighlight = action
+        self.console.print(f"Set highlight to {self.vhighlight}")
+
+    def limit(self, command):
+        """
+        Processes .limit command.
+
+        Args:
+            command: command line
+        """
+
+        _, action = self.split(command, 10)
+        self.vlimit = int(action)
+        self.console.print(f"Set limit to {self.vlimit}")
+
+    def load(self, path):
+        """
+        Processes .load command.
+
+        Args:
+            path: path to configuration
+        """
+
+        if self.isyaml(path):
+            self.console.print(f"Loading application {path}")
+            self.app = Application(path)
+        else:
+            self.console.print(f"Loading index {path}")
+
+            # Load embeddings index
+            self.app = Embeddings()
+            self.app.load(path)
+
+    def search(self, query):
+        """
+        Runs a search query.
+
+        Args:
+            query: query to run
+        """
+
+        if self.vhighlight:
+            results = self.app.explain(query, limit=self.vlimit)
+        else:
+            results = self.app.search(query, limit=self.vlimit)
+
+        columns, table = {}, Table(box=box.SQUARE, style="#03a9f4")
+
+        # Build column list
+        result = results[0]
+        if isinstance(result, tuple):
+            columns = dict.fromkeys(["id", "score"])
+        else:
+            columns = dict(result)
+
+        # Add columns to table
+        columns = list(x for x in columns if x != "tokens")
+        for column in columns:
+            table.add_column(column)
+
+        # Add rows to table
+        for result in results:
+            if isinstance(result, tuple):
+                table.add_row(*(self.render(result, None, x) for x in result))
+            else:
+                table.add_row(*(self.render(result, column, result.get(column)) for column in columns))
+
+        # Print table to console
+        self.console.print(table)
+
+    def workflow(self, command):
+        """
+        Processes .workflow command.
+
+        Args:
+            command: command line
+        """
+
+        command = shlex.split(command)
+        if isinstance(self.app, Application):
+            self.console.print(list(self.app.workflow(command[1], command[2:])))
+
+    def isyaml(self, path):
+        """
+        Checks if file at path is a valid YAML file.
+
+        Args:
+            path: file to check
+
+        Returns:
+            True if file is valid YAML, False otherwise
+        """
+
+        if os.path.exists(path) and os.path.isfile(path):
+            try:
+                return Application.read(path)
+            # pylint: disable=W0702
+            except:
+                pass
+
+        return False
+
+    def split(self, command, default=None):
+        """
+        Splits command by whitespace.
+
+        Args:
+            command: command line
+            default: default command action
+
+        Returns:
+            command action
+        """
+
+        values = command.split(" ", 1)
+        return values if len(values) > 1 else (command, default)
+
+    def render(self, result, column, value):
+        """
+        Renders a search result column value.
+
+        Args:
+            result: result row
+            column: column name
+            value: column value
+        """
+
+        if isinstance(value, float):
+            return f"{value:.4f}"
+
+        # Explain highlighting
+        if column == "text" and "tokens" in result:
+            spans = []
+            for token, score in result["tokens"]:
+                color = None
+                if score >= 0.02:
+                    color = f"b {self.vhighlight}"
+
+                spans.append((token, score, color))
+
+            if result["score"] >= 0.05 and not [color for _, _, color in spans if color]:
+                mscore = max(score for _, score, _ in spans)
+                spans = [(token, score, f"b {self.vhighlight}" if score == mscore else color) for token, score, color in spans]
+
+            output = ""
+            for token, _, color in spans:
+                if color:
+                    output += f"[{color}]{token}[/{color}] "
+                else:
+                    output += f"{token} "
+
+            return output
+
+        return str(value)

txtai/data/__init__.py

@@ -0,0 +1,10 @@
+"""
+Data imports
+"""
+
+from .base import Data
+from .labels import Labels
+from .questions import Questions
+from .sequences import Sequences
+from .texts import Texts
+from .tokens import Tokens

txtai/data/base.py

@@ -0,0 +1,138 @@
+"""
+Data module
+"""
+
+from .tokens import Tokens
+
+
+class Data:
+    """
+    Base data tokenization class.
+    """
+
+    def __init__(self, tokenizer, columns, maxlength):
+        """
+        Creates new base instance for tokenizing data.
+
+        Args:
+            tokenizer: model tokenizer
+            columns: column names
+            maxlength: maximum sequence length
+        """
+
+        self.tokenizer = tokenizer
+        self.columns = columns
+        self.maxlength = maxlength
+
+    def __call__(self, train, validation, workers):
+        """
+        Tokenizes training and validation data and returns processed datasets.
+
+        Args:
+            train: training data
+            validation: validation data
+            workers: number of concurrent tokenizers when processing datasets, only main process used when set to None
+
+        Returns:
+            (train, validation)
+        """
+
+        return (self.prepare(train, self.process, workers), self.prepare(validation, self.process, workers) if validation else None)
+
+    def prepare(self, data, fn, workers):
+        """
+        Prepares and tokenizes data for training.
+
+        Args:
+            data: input data
+            fn: tokenize processing function to apply
+            workers: number of concurrent tokenizers when processing datasets, only main process used when set to None
+
+        Returns:
+            tokens
+        """
+
+        if hasattr(data, "map"):
+            # Hugging Face dataset
+            tokens = data.map(fn, batched=True, num_proc=workers, remove_columns=data.column_names)
+        else:
+            # Re-orient data into columns for efficient batch tokenization
+            columns = {}
+            if hasattr(data, "columns"):
+                # Polars/pandas DataFrame
+                for column in data.columns:
+                    columns[column] = list(data[column])
+            else:
+                # Iterable dicts
+                for row in data:
+                    for column in row.keys():
+                        if column not in columns:
+                            columns[column] = []
+
+                        columns[column].append(row[column])
+
+            # Process column-oriented data
+            tokens = Tokens(fn(columns))
+
+        return tokens
+
+    def labels(self, data):
+        """
+        Extracts a list of unique labels from data.
+
+        Args:
+            data: input data
+
+        Returns:
+            list of unique labels
+        """
+
+        # Last column is label
+        column = self.columns[-1]
+
+        # Return length of labels if it's an array
+        length = self.length(data[column][0] if hasattr(data, "columns") else data[0][column])
+        if length:
+            return length
+
+        if hasattr(data, "map"):
+            # Hugging Face dataset
+            labels = sorted(data.unique(self.columns[-1]))
+        elif hasattr(data, "columns"):
+            # Polars/pandas DataFrame
+            labels = sorted(data[self.columns[-1]].unique())
+        else:
+            # Iterable dicts
+            labels = sorted({row[self.columns[-1]] for row in data})
+
+        # Labels are single numeric values per entry
+        #   - Consider a regression task if at least one label isn't an integer
+        #   - Otherwise use number of labels for a classification task
+        return 1 if [x for x in labels if float(x) != int(x)] else len(labels)
+
+    def process(self, data):
+        """
+        Tokenizes batch of input data
+
+        Args:
+            data: input data batch
+
+        Returns:
+            tokenized data
+        """
+
+        return data
+
+    def length(self, value):
+        """
+        Returns the length of value if value has a len function defined. Otherwise,
+        None is returned.
+
+        Args:
+            value: value to check
+
+        Returns:
+            length of value if available, otherwise returns None
+        """
+
+        return len(value) if hasattr(value, "__len__") else None

txtai/data/labels.py

@@ -0,0 +1,42 @@
+"""
+Labels module
+"""
+
+from .base import Data
+
+
+class Labels(Data):
+    """
+    Tokenizes text-classification datasets as input for training text-classification models.
+    """
+
+    def __init__(self, tokenizer, columns, maxlength):
+        """
+        Creates a new instance for tokenizing Labels training data.
+
+        Args:
+            tokenizer: model tokenizer
+            columns: tuple of columns to use for text/label
+            maxlength: maximum sequence length
+        """
+
+        super().__init__(tokenizer, columns, maxlength)
+
+        # Standardize columns
+        if not self.columns:
+            self.columns = ("text", None, "label")
+        elif len(columns) < 3:
+            self.columns = (self.columns[0], None, self.columns[-1])
+
+    def process(self, data):
+        # Column keys
+        text1, text2, label = self.columns
+
+        # Tokenizer inputs can be single string or string pair, depending on task
+        text = (data[text1], data[text2]) if text2 else (data[text1],)
+
+        # Tokenize text and add label
+        inputs = self.tokenizer(*text, max_length=self.maxlength, padding=True, truncation=True)
+        inputs[label] = data[label]
+
+        return inputs

txtai/data/questions.py

@@ -0,0 +1,135 @@
+"""
+Questions module
+"""
+
+from .base import Data
+
+
+class Questions(Data):
+    """
+    Tokenizes question-answering datasets as input for training question-answering models.
+    """
+
+    def __init__(self, tokenizer, columns, maxlength, stride):
+        """
+        Creates a new instance for tokenizing Questions training data.
+
+        Args:
+            tokenizer: model tokenizer
+            columns: tuple of columns to use for question/context/answer
+            maxlength: maximum sequence length
+            stride: chunk size for splitting data for QA tasks
+        """
+
+        super().__init__(tokenizer, columns, maxlength)
+
+        if not self.columns:
+            self.columns = ("question", "context", "answers")
+
+        self.question, self.context, self.answer = self.columns
+        self.stride = stride
+        self.rpad = tokenizer.padding_side == "right"
+
+    def process(self, data):
+        # Tokenize data
+        tokenized = self.tokenize(data)
+
+        # Get mapping of overflowing tokens and answer offsets
+        samples = tokenized.pop("overflow_to_sample_mapping")
+        offsets = tokenized.pop("offset_mapping")
+
+        # Start/end positions
+        tokenized["start_positions"] = []
+        tokenized["end_positions"] = []
+
+        for x, offset in enumerate(offsets):
+            # Label NO ANSWER with CLS token
+            inputids = tokenized["input_ids"][x]
+            clstoken = inputids.index(self.tokenizer.cls_token_id)
+
+            # Sequence ids
+            sequences = tokenized.sequence_ids(x)
+
+            # Get and format answer
+            answers = self.answers(data, samples[x])
+
+            # If no answers are given, set cls token as answer.
+            if len(answers["answer_start"]) == 0:
+                tokenized["start_positions"].append(clstoken)
+                tokenized["end_positions"].append(clstoken)
+            else:
+                # Start/end character index of the answer in the text.
+                startchar = answers["answer_start"][0]
+                endchar = startchar + len(answers["text"][0])
+
+                # Start token index of the current span in the text.
+                start = 0
+                while sequences[start] != (1 if self.rpad else 0):
+                    start += 1
+
+                # End token index of the current span in the text.
+                end = len(inputids) - 1
+                while sequences[end] != (1 if self.rpad else 0):
+                    end -= 1
+
+                # Map start character and end character to matching token index
+                while start < len(offset) and offset[start][0] <= startchar:
+                    start += 1
+                tokenized["start_positions"].append(start - 1)
+
+                while offset[end][1] >= endchar:
+                    end -= 1
+                tokenized["end_positions"].append(end + 1)
+
+        return tokenized
+
+    def tokenize(self, data):
+        """
+        Tokenizes batch of data
+
+        Args:
+            data: input data batch
+
+        Returns:
+            tokenized data
+        """
+
+        # Trim question whitespace
+        data[self.question] = [x.lstrip() for x in data[self.question]]
+
+        # Tokenize records
+        return self.tokenizer(
+            data[self.question if self.rpad else self.context],
+            data[self.context if self.rpad else self.question],
+            truncation="only_second" if self.rpad else "only_first",
+            max_length=self.maxlength,
+            stride=self.stride,
+            return_overflowing_tokens=True,
+            return_offsets_mapping=True,
+            padding=True,
+        )
+
+    def answers(self, data, index):
+        """
+        Gets and formats an answer.
+
+        Args:
+            data: input examples
+            index: answer index to retrieve
+
+        Returns:
+            answers dict
+        """
+
+        # Answer mappings
+        answers = data[self.answer][index]
+        context = data[self.context][index]
+
+        # Handle mapping string answers to dict
+        if not isinstance(answers, dict):
+            if not answers:
+                answers = {"text": [], "answer_start": []}
+            else:
+                answers = {"text": [answers], "answer_start": [context.index(answers)]}
+
+        return answers

txtai/data/sequences.py

@@ -0,0 +1,48 @@
+"""
+Sequences module
+"""
+
+from .base import Data
+
+
+class Sequences(Data):
+    """
+    Tokenizes sequence-sequence datasets as input for training sequence-sequence models
+    """
+
+    def __init__(self, tokenizer, columns, maxlength, prefix):
+        """
+        Creates a new instance for tokenizing Sequences training data.
+
+        Args:
+            tokenizer: model tokenizer
+            columns: tuple of columns to use for text/label
+            maxlength: maximum sequence length
+            prefix: source prefix
+        """
+
+        super().__init__(tokenizer, columns, maxlength)
+
+        # Standardize columns
+        if not self.columns:
+            self.columns = ("source", "target")
+
+        # Save source prefix
+        self.prefix = prefix
+
+    def process(self, data):
+        # Column keys
+        source, target = self.columns
+
+        # Tokenize source
+        source = [self.prefix + x if self.prefix else x for x in data[source]]
+        inputs = self.tokenizer(source, max_length=self.maxlength, padding=False, truncation=True)
+
+        # Tokenize target
+        with self.tokenizer.as_target_tokenizer():
+            targets = self.tokenizer(data[target], max_length=self.maxlength, padding=False, truncation=True)
+
+        # Combine inputs
+        inputs["labels"] = targets["input_ids"]
+
+        return inputs