mseep-txtai 9.1.1__py3-none-any.whl

txtai/data/texts.py

@@ -0,0 +1,68 @@
+"""
+Texts module
+"""
+
+from itertools import chain
+
+from .base import Data
+
+
+class Texts(Data):
+    """
+    Tokenizes text datasets as input for training language models.
+    """
+
+    def __init__(self, tokenizer, columns, maxlength):
+        """
+        Creates a new instance for tokenizing Texts training data.
+
+        Args:
+            tokenizer: model tokenizer
+            columns: tuple of columns to use for text
+            maxlength: maximum sequence length
+        """
+
+        super().__init__(tokenizer, columns, maxlength)
+
+        # Standardize columns
+        if not self.columns:
+            self.columns = ("text", None)
+
+    def process(self, data):
+        # Column keys
+        text1, text2 = 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, return_special_tokens_mask=True)
+
+        # Concat and return tokenized inputs
+        return self.concat(inputs)
+
+    def concat(self, inputs):
+        """
+        Concatenates tokenized text into chunks of maxlength.
+
+        Args:
+            inputs: tokenized input
+
+        Returns:
+            Chunks of tokenized text each with a size of maxlength
+        """
+
+        # Concatenate tokenized text
+        concat = {k: list(chain(*inputs[k])) for k in inputs.keys()}
+
+        # Calculate total length
+        length = len(concat[list(inputs.keys())[0]])
+
+        # Ensure total is multiple of maxlength, drop last incomplete chunk
+        if length >= self.maxlength:
+            length = (length // self.maxlength) * self.maxlength
+
+        # Split into chunks of maxlength
+        result = {k: [v[x : x + self.maxlength] for x in range(0, length, self.maxlength)] for k, v in concat.items()}
+
+        return result

txtai/data/tokens.py

@@ -0,0 +1,28 @@
+"""
+Tokens module
+"""
+
+import torch
+
+
+class Tokens(torch.utils.data.Dataset):
+    """
+    Default dataset used to hold tokenized data.
+    """
+
+    def __init__(self, columns):
+        self.data = []
+
+        # Map column-oriented data to rows
+        for column in columns:
+            for x, value in enumerate(columns[column]):
+                if len(self.data) <= x:
+                    self.data.append({})
+
+                self.data[x][column] = value
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]

txtai/database/__init__.py

@@ -0,0 +1,14 @@
+"""
+Database imports
+"""
+
+from .base import Database
+from .client import Client
+from .duckdb import DuckDB
+from .embedded import Embedded
+from .encoder import *
+from .factory import DatabaseFactory
+from .rdbms import RDBMS
+from .schema import *
+from .sqlite import SQLite
+from .sql import *

txtai/database/base.py

@@ -0,0 +1,342 @@
+"""
+Database module
+"""
+
+import logging
+import types
+
+from .encoder import EncoderFactory
+from .sql import SQL, SQLError, Token
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+
+class Database:
+    """
+    Base class for database instances. This class encapsulates a content database used for
+    storing field content as dicts and objects. The database instance works in conjuction
+    with a vector index to execute SQL-driven similarity search.
+    """
+
+    def __init__(self, config):
+        """
+        Creates a new Database.
+
+        Args:
+            config: database configuration
+        """
+
+        # Initialize configuration
+        self.configure(config)
+
+    def load(self, path):
+        """
+        Loads a database path.
+
+        Args:
+            path: database url
+        """
+
+        raise NotImplementedError
+
+    def insert(self, documents, index=0):
+        """
+        Inserts documents into the database.
+
+        Args:
+            documents: list of documents to save
+            index: indexid offset, used for internal ids
+        """
+
+        raise NotImplementedError
+
+    def delete(self, ids):
+        """
+        Deletes documents from database.
+
+        Args:
+            ids: ids to delete
+        """
+
+        raise NotImplementedError
+
+    def reindex(self, config):
+        """
+        Reindexes internal database content and streams results back. This method must renumber indexids
+        sequentially as deletes could have caused indexid gaps.
+
+        Args:
+            config: new configuration
+        """
+
+        raise NotImplementedError
+
+    def save(self, path):
+        """
+        Saves a database at path.
+
+        Args:
+            path: path to write database
+        """
+
+        raise NotImplementedError
+
+    def close(self):
+        """
+        Closes this database.
+        """
+
+        raise NotImplementedError
+
+    def ids(self, ids):
+        """
+        Retrieves the internal indexids for a list of ids. Multiple indexids may be present for an id in cases
+        where data is segmented.
+
+        Args:
+            ids: list of document ids
+
+        Returns:
+            list of (indexid, id)
+        """
+
+        raise NotImplementedError
+
+    def count(self):
+        """
+        Retrieves the count of this database instance.
+
+        Returns:
+            total database count
+        """
+
+        raise NotImplementedError
+
+    def search(self, query, similarity=None, limit=None, parameters=None, indexids=False):
+        """
+        Runs a search against the database. Supports the following methods:
+
+            1. Standard similarity query. This mode retrieves content for the ids in the similarity results
+            2. Similarity query as SQL. This mode will combine similarity results and database results into
+               a single result set. Similarity queries are set via the SIMILAR() function.
+            3. SQL with no similarity query. This mode runs a SQL query and retrieves the results without similarity queries.
+
+        Example queries:
+            "natural language processing" - standard similarity only query
+            "select * from txtai where similar('natural language processing')" - similarity query as SQL
+            "select * from txtai where similar('nlp') and entry > '2021-01-01'" - similarity query with additional SQL clauses
+            "select id, text, score from txtai where similar('nlp')" - similarity query with additional SQL column selections
+            "select * from txtai where entry > '2021-01-01' - database only query
+
+        Args:
+            query: input query
+            similarity: similarity results as [(indexid, score)]
+            limit: maximum number of results to return
+            parameters: dict of named parameters to bind to placeholders
+
+        Returns:
+            query results as a list of dicts
+            list of ([indexid, score]) if indexids is True
+        """
+
+        # Parse query if necessary
+        if isinstance(query, str):
+            query = self.parse(query)
+
+        # Add in similar results
+        where = query.get("where")
+
+        if "select" in query and similarity:
+            for x in range(len(similarity)):
+                token = f"{Token.SIMILAR_TOKEN}{x}"
+                if where and token in where:
+                    where = where.replace(token, self.embed(similarity, x))
+
+        elif similarity:
+            # Not a SQL query, load similarity results, if any
+            where = self.embed(similarity, 0)
+
+        # Save where
+        query["where"] = where
+
+        # Run query
+        return self.query(query, limit, parameters, indexids)
+
+    def parse(self, query):
+        """
+        Parses a query into query components.
+
+        Args:
+            query: input query
+
+        Returns:
+            dict of parsed query components
+        """
+
+        return self.sql(query)
+
+    def resolve(self, name, alias=None):
+        """
+        Resolves a query column name with the database column name. This method also builds alias expressions
+        if alias is set.
+
+        Args:
+            name: query column name
+            alias: alias name, defaults to None
+
+        Returns:
+            database column name
+        """
+
+        raise NotImplementedError
+
+    def embed(self, similarity, batch):
+        """
+        Embeds similarity query results into a database query.
+
+        Args:
+            similarity: similarity results as [(indexid, score)]
+            batch: batch id
+        """
+
+        raise NotImplementedError
+
+    def query(self, query, limit, parameters, indexids):
+        """
+        Executes query against database.
+
+        Args:
+            query: input query
+            limit: maximum number of results to return
+            parameters: dict of named parameters to bind to placeholders
+            indexids: results are returned as [(indexid, score)] regardless of select clause parameters if True
+
+        Returns:
+            query results
+        """
+
+        raise NotImplementedError
+
+    def configure(self, config):
+        """
+        Initialize configuration.
+
+        Args:
+            config: configuration
+        """
+
+        # Database configuration
+        self.config = config
+
+        # SQL parser
+        self.sql = SQL(self)
+
+        # Load objects encoder
+        encoder = self.config.get("objects")
+        self.encoder = EncoderFactory.create(encoder) if encoder else None
+
+        # Transform columns
+        columns = config.get("columns", {})
+        self.text = columns.get("text", "text")
+        self.object = columns.get("object", "object")
+
+        # Custom functions and expressions
+        self.functions, self.expressions = None, None
+
+        # Load custom functions
+        self.registerfunctions(self.config)
+
+        # Load custom expressions
+        self.registerexpressions(self.config)
+
+    def registerfunctions(self, config):
+        """
+        Register custom functions. This method stores the function details for underlying
+        database implementations to handle.
+
+        Args:
+            config: database configuration
+        """
+
+        inputs = config.get("functions") if config else None
+        if inputs:
+            functions = []
+            for fn in inputs:
+                name, argcount = None, -1
+
+                # Optional function configuration
+                if isinstance(fn, dict):
+                    name, argcount, fn = fn.get("name"), fn.get("argcount", -1), fn["function"]
+
+                # Determine if this is a callable object or a function
+                if not isinstance(fn, types.FunctionType) and hasattr(fn, "__call__"):
+                    name = name if name else fn.__class__.__name__.lower()
+                    fn = fn.__call__
+                else:
+                    name = name if name else fn.__name__.lower()
+
+                # Store function details
+                functions.append((name, argcount, fn))
+
+            # pylint: disable=W0201
+            self.functions = functions
+
+    def registerexpressions(self, config):
+        """
+        Register custom expressions. This method parses and resolves expressions for later use in SQL queries.
+
+        Args:
+            config: database configuration
+        """
+
+        inputs = config.get("expressions") if config else None
+        if inputs:
+            expressions = {}
+            for entry in inputs:
+                name = entry.get("name")
+                expression = entry.get("expression")
+                if name and expression:
+                    expressions[name] = self.sql.snippet(expression)
+
+            # pylint: disable=W0201
+            self.expressions = expressions
+
+    def execute(self, function, *args):
+        """
+        Executes a user query. This method has common error handling logic.
+
+        Args:
+            function: database execute function
+            args: function arguments
+
+        Returns:
+            result of function(args)
+        """
+
+        try:
+            # Debug log SQL
+            logger.debug(" ".join(["%s"] * len(args)), *args)
+
+            return function(*args)
+        except Exception as e:
+            raise SQLError(e) from None
+
+    def setting(self, name, default=None):
+        """
+        Looks up database specific setting.
+
+        Args:
+            name: setting name
+            default: default value when setting not found
+
+        Returns:
+            setting value
+        """
+
+        # Get the database-specific config object
+        database = self.config.get(self.config["content"])
+
+        # Get setting value, set default value if not found
+        setting = database.get(name) if database else None
+        return setting if setting else default

txtai/database/client.py

@@ -0,0 +1,227 @@
+"""
+Client module
+"""
+
+import os
+import time
+
+# Conditional import
+try:
+    from sqlalchemy import StaticPool, Text, cast, create_engine, insert, text as textsql
+    from sqlalchemy.orm import Session, aliased
+    from sqlalchemy.schema import CreateSchema
+
+    from .schema import Base, Batch, Document, Object, Section, SectionBase, Score
+
+    ORM = True
+except ImportError:
+    ORM = False
+
+from .rdbms import RDBMS
+
+
+class Client(RDBMS):
+    """
+    Database client instance. This class connects to an external database using SQLAlchemy. It supports any database
+    that is supported by SQLAlchemy (PostgreSQL, MariaDB, etc) and has JSON support.
+    """
+
+    def __init__(self, config):
+        """
+        Creates a new Database.
+
+        Args:
+            config: database configuration parameters
+        """
+
+        super().__init__(config)
+
+        if not ORM:
+            raise ImportError('SQLAlchemy is not available - install "database" extra to enable')
+
+        # SQLAlchemy parameters
+        self.engine, self.dbconnection = None, None
+
+    def save(self, path):
+        # Commit session and database connection
+        super().save(path)
+
+        if self.dbconnection:
+            self.dbconnection.commit()
+
+    def close(self):
+        super().close()
+
+        # Dispose of engine, which also closes dbconnection
+        if self.engine:
+            self.engine.dispose()
+
+    def reindexstart(self):
+        # Working table name
+        name = f"rebuild{round(time.time() * 1000)}"
+
+        # Create working table metadata
+        type("Rebuild", (SectionBase,), {"__tablename__": name})
+        Base.metadata.tables[name].create(self.dbconnection)
+
+        return name
+
+    def reindexend(self, name):
+        # Remove table object from metadata
+        Base.metadata.remove(Base.metadata.tables[name])
+
+    def jsonprefix(self):
+        # JSON column prefix
+        return "cast("
+
+    def jsoncolumn(self, name):
+        # Alias documents table
+        d = aliased(Document, name="d")
+
+        # Build JSON column expression for column
+        return str(cast(d.data[name].as_string(), Text).compile(dialect=self.engine.dialect, compile_kwargs={"literal_binds": True}))
+
+    def createtables(self):
+        # Create tables
+        Base.metadata.create_all(self.dbconnection, checkfirst=True)
+
+        # Clear existing data - table schema is created upon connecting to database
+        for table in ["sections", "documents", "objects"]:
+            self.cursor.execute(f"DELETE FROM {table}")
+
+    def finalize(self):
+        # Flush cached objects
+        self.connection.flush()
+
+    def insertdocument(self, uid, data, tags, entry):
+        self.connection.add(Document(id=uid, data=data, tags=tags, entry=entry))
+
+    def insertobject(self, uid, data, tags, entry):
+        self.connection.add(Object(id=uid, object=data, tags=tags, entry=entry))
+
+    def insertsection(self, index, uid, text, tags, entry):
+        # Save text section
+        self.connection.add(Section(indexid=index, id=uid, text=text, tags=tags, entry=entry))
+
+    def createbatch(self):
+        # Create temporary batch table, if necessary
+        Base.metadata.tables["batch"].create(self.dbconnection, checkfirst=True)
+
+    def insertbatch(self, indexids, ids, batch):
+        if indexids:
+            self.connection.execute(insert(Batch), [{"indexid": i, "batch": batch} for i in indexids])
+        if ids:
+            self.connection.execute(insert(Batch), [{"id": str(uid), "batch": batch} for uid in ids])
+
+    def createscores(self):
+        # Create temporary scores table, if necessary
+        Base.metadata.tables["scores"].create(self.dbconnection, checkfirst=True)
+
+    def insertscores(self, scores):
+        # Average scores by id
+        if scores:
+            self.connection.execute(insert(Score), [{"indexid": i, "score": sum(s) / len(s)} for i, s in scores.items()])
+
+    def connect(self, path=None):
+        # Connection URL
+        content = self.config.get("content")
+
+        # Read ENV variable, if necessary
+        content = os.environ.get("CLIENT_URL") if content == "client" else content
+
+        # Create engine using database URL
+        self.engine = create_engine(content, poolclass=StaticPool, echo=False, json_serializer=lambda x: x)
+        self.dbconnection = self.engine.connect()
+
+        # Create database session
+        database = Session(self.dbconnection)
+
+        # Set default schema, if necessary
+        schema = self.config.get("schema")
+        if schema:
+            with self.engine.begin():
+                self.sqldialect(database, CreateSchema(schema, if_not_exists=True))
+
+            self.sqldialect(database, textsql("SET search_path TO :schema"), {"schema": schema})
+
+        return database
+
+    def getcursor(self):
+        return Cursor(self.connection)
+
+    def rows(self):
+        return self.cursor
+
+    def addfunctions(self):
+        return
+
+    def sqldialect(self, database, sql, parameters=None):
+        """
+        Executes a SQL statement based on the current SQL dialect.
+
+        Args:
+            database: current database
+            sql: SQL to execute
+            parameters: optional bind parameters
+        """
+
+        args = (sql, parameters) if self.engine.dialect.name == "postgresql" else (textsql("SELECT 1"),)
+        database.execute(*args)
+
+
+class Cursor:
+    """
+    Implements basic compatibility with the Python DB-API.
+    """
+
+    def __init__(self, connection):
+        self.connection = connection
+        self.result = None
+
+    def __iter__(self):
+        return self.result
+
+    def execute(self, statement, parameters=None):
+        """
+        Executes statement.
+
+        Args:
+            statement: statement to execute
+            parameters: optional dictionary with bind parameters
+        """
+
+        if isinstance(statement, str):
+            statement = textsql(statement)
+
+        self.result = self.connection.execute(statement, parameters)
+
+    def fetchall(self):
+        """
+        Fetches all rows from the current result.
+
+        Returns:
+            all rows from current result
+        """
+
+        return self.result.all() if self.result else None
+
+    def fetchone(self):
+        """
+        Fetches first row from current result.
+
+        Returns:
+            first row from current result
+        """
+
+        return self.result.first() if self.result else None
+
+    @property
+    def description(self):
+        """
+        Returns columns for current result.
+
+        Returns:
+            list of columns
+        """
+
+        return [(key,) for key in self.result.keys()] if self.result else None