mseep-txtai 9.1.1__py3-none-any.whl

txtai/archive/__init__.py

@@ -0,0 +1,9 @@
+"""
+Archive imports
+"""
+
+from .base import Archive
+from .compress import Compress
+from .factory import ArchiveFactory
+from .tar import Tar
+from .zip import Zip

txtai/archive/base.py

@@ -0,0 +1,104 @@
+"""
+Archive module
+"""
+
+import os
+
+from tempfile import TemporaryDirectory
+
+from .tar import Tar
+from .zip import Zip
+
+
+class Archive:
+    """
+    Base class for archive instances.
+    """
+
+    def __init__(self, directory=None):
+        """
+        Creates a new archive instance.
+
+        Args:
+            directory: directory to use as working directory, defaults to a temporary directory
+        """
+
+        self.directory = directory
+
+    def isarchive(self, path):
+        """
+        Checks if path is an archive file based on the extension.
+
+        Args:
+            path: path to check
+
+        Returns:
+            True if the path ends with an archive extension, False otherwise
+        """
+
+        return path and any(path.lower().endswith(extension) for extension in [".tar.bz2", ".tar.gz", ".tar.xz", ".zip"])
+
+    def path(self):
+        """
+        Gets the current working directory for this archive instance.
+
+        Returns:
+            archive working directory
+        """
+
+        # Default to a temporary directory. All files created in this directory will be deleted
+        # when this archive instance goes out of scope.
+        if not self.directory:
+            # pylint: disable=R1732
+            self.directory = TemporaryDirectory()
+
+        return self.directory.name if isinstance(self.directory, TemporaryDirectory) else self.directory
+
+    def load(self, path, compression=None):
+        """
+        Extracts file at path to archive working directory.
+
+        Args:
+            path: path to archive file
+            compression: compression format, infers from path if not provided
+        """
+
+        # Unpack compressed file
+        compress = self.create(path, compression)
+        compress.unpack(path, self.path())
+
+    def save(self, path, compression=None):
+        """
+        Archives files in archive working directory to file at path.
+
+        Args:
+            path: path to archive file
+            compression: compression format, infers from path if not provided
+        """
+
+        # Create output directory, if necessary
+        output = os.path.dirname(path)
+        if output:
+            os.makedirs(output, exist_ok=True)
+
+        # Pack into compressed file
+        compress = self.create(path, compression)
+        compress.pack(self.path(), path)
+
+    def create(self, path, compression):
+        """
+        Method to construct a Compress instance.
+
+        Args:
+            path: file path
+            compression: compression format, infers using file extension if not provided
+
+        Returns:
+            Compress
+        """
+
+        # Infer compression format from path if not provided
+        compression = compression if compression else path.lower().split(".")[-1]
+
+        # Create compression instance
+        return Zip() if compression == "zip" else Tar()

txtai/archive/compress.py

@@ -0,0 +1,51 @@
+"""
+Compress module
+"""
+
+import os
+
+
+class Compress:
+    """
+    Base class for Compress instances.
+    """
+
+    def pack(self, path, output):
+        """
+        Compresses files in directory path to file output.
+
+        Args:
+            path: input directory path
+            output: output file
+        """
+
+        raise NotImplementedError
+
+    def unpack(self, path, output):
+        """
+        Extracts all files in path to output.
+
+        Args:
+            path: input file path
+            output: output directory
+        """
+
+        raise NotImplementedError
+
+    def validate(self, directory, path):
+        """
+        Validates path is under directory.
+
+        Args:
+            directory: base directory
+            path: path to validate
+
+        Returns:
+            True if path is under directory, False otherwise
+        """
+
+        directory = os.path.abspath(directory)
+        path = os.path.abspath(path)
+        prefix = os.path.commonprefix([directory, path])
+
+        return prefix == directory

txtai/archive/factory.py

@@ -0,0 +1,25 @@
+"""
+Factory module
+"""
+
+from .base import Archive
+
+
+class ArchiveFactory:
+    """
+    Methods to create Archive instances.
+    """
+
+    @staticmethod
+    def create(directory=None):
+        """
+        Create a new Archive instance.
+
+        Args:
+            directory: optional default working directory, otherwise uses a temporary directory
+
+        Returns:
+            Archive
+        """
+
+        return Archive(directory)

txtai/archive/tar.py

@@ -0,0 +1,49 @@
+"""
+Tar module
+"""
+
+import os
+import tarfile
+
+from .compress import Compress
+
+
+class Tar(Compress):
+    """
+    Tar compression
+    """
+
+    def pack(self, path, output):
+        # Infer compression type
+        compression = self.compression(output)
+
+        with tarfile.open(output, f"w:{compression}" if compression else "w") as tar:
+            tar.add(path, arcname=".")
+
+    def unpack(self, path, output):
+        # Infer compression type
+        compression = self.compression(path)
+
+        with tarfile.open(path, f"r:{compression}" if compression else "r") as tar:
+            # Validate paths
+            for member in tar.getmembers():
+                fullpath = os.path.join(path, member.name)
+                if not self.validate(path, fullpath):
+                    raise IOError(f"Invalid tar entry: {member.name}")
+
+            tar.extractall(output)
+
+    def compression(self, path):
+        """
+        Gets compression type for path.
+
+        Args:
+            path: path to file
+
+        Returns:
+            compression type
+        """
+
+        # Infer compression type from last path component. Limit to supported types.
+        compression = path.lower().split(".")[-1]
+        return compression if compression in ("bz2", "gz", "xz") else None

txtai/archive/zip.py

@@ -0,0 +1,35 @@
+"""
+Zip module
+"""
+
+import os
+
+from zipfile import ZipFile, ZIP_DEFLATED
+
+from .compress import Compress
+
+
+class Zip(Compress):
+    """
+    Zip compression
+    """
+
+    def pack(self, path, output):
+        with ZipFile(output, "w", ZIP_DEFLATED) as zfile:
+            for root, _, files in sorted(os.walk(path)):
+                for f in files:
+                    # Generate archive name with relative path, if necessary
+                    name = os.path.join(os.path.relpath(root, path), f)
+
+                    # Write file to zip
+                    zfile.write(os.path.join(root, f), arcname=name)
+
+    def unpack(self, path, output):
+        with ZipFile(path, "r") as zfile:
+            # Validate path if directory specified
+            for fullpath in zfile.namelist():
+                fullpath = os.path.join(path, fullpath)
+                if os.path.dirname(fullpath) and not self.validate(path, fullpath):
+                    raise IOError(f"Invalid zip entry: {fullpath}")
+
+            zfile.extractall(output)

txtai/cloud/__init__.py

@@ -0,0 +1,8 @@
+"""
+Cloud imports
+"""
+
+from .base import Cloud
+from .factory import CloudFactory
+from .hub import HuggingFaceHub
+from .storage import ObjectStorage

txtai/cloud/base.py

@@ -0,0 +1,106 @@
+"""
+Cloud module
+"""
+
+import os
+
+from ..archive import ArchiveFactory
+
+
+class Cloud:
+    """
+    Base class for cloud providers. Cloud providers sync content between local and remote storage.
+    """
+
+    def __init__(self, config):
+        """
+        Creates a new cloud connection.
+
+        Args:
+            config: cloud configuration
+        """
+
+        self.config = config
+
+    def exists(self, path=None):
+        """
+        Checks if path exists in cloud. If path is None, this method checks if the container exists.
+
+        Args:
+            path: path to check
+
+        Returns:
+            True if path or container exists, False otherwise
+        """
+
+        return self.metadata(path) is not None
+
+    def metadata(self, path=None):
+        """
+        Returns metadata for path from cloud. If path is None, this method returns metadata
+        for container.
+
+        Args:
+            path: retrieve metadata for this path
+
+        Returns:
+            path or container metadata if available, otherwise returns None
+        """
+
+        raise NotImplementedError
+
+    def load(self, path=None):
+        """
+        Retrieves content from cloud and stores locally. If path is empty, this method retrieves
+        all content in the container.
+
+        Args:
+            path: path to retrieve
+
+        Returns:
+            local path which can be different than input path
+        """
+
+        raise NotImplementedError
+
+    def save(self, path):
+        """
+        Sends local content stored in path to cloud.
+
+        Args:
+            path: local path to sync
+        """
+
+        raise NotImplementedError
+
+    def isarchive(self, path):
+        """
+        Check if path is an archive file.
+
+        Args:
+            path: path to check
+
+        Returns:
+            True if path ends with an archive extension, false otherwise
+        """
+
+        return ArchiveFactory.create().isarchive(path)
+
+    def listfiles(self, path):
+        """
+        Lists files in path. If path is a file, this method returns a single element list
+        containing path.
+
+        Args:
+            path: path to list
+
+        Returns:
+            List of files
+        """
+
+        # List all files if path is a directory
+        if os.path.isdir(path):
+            return [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
+
+        # Path is a file
+        return [path]

txtai/cloud/factory.py

@@ -0,0 +1,70 @@
+"""
+Factory module
+"""
+
+from ..util import Resolver
+
+from .hub import HuggingFaceHub
+from .storage import ObjectStorage, LIBCLOUD
+
+
+class CloudFactory:
+    """
+    Methods to create Cloud instances.
+    """
+
+    @staticmethod
+    def create(config):
+        """
+        Creates a Cloud instance.
+
+        Args:
+            config: cloud configuration
+
+        Returns:
+            Cloud
+        """
+
+        # Cloud instance
+        cloud = None
+
+        provider = config.get("provider", "")
+
+        # Hugging Face Hub
+        if provider.lower() == "huggingface-hub":
+            cloud = HuggingFaceHub(config)
+
+        # Cloud object storage
+        elif ObjectStorage.isprovider(provider):
+            cloud = ObjectStorage(config)
+
+        # External provider
+        elif provider:
+            cloud = CloudFactory.resolve(provider, config)
+
+        return cloud
+
+    @staticmethod
+    def resolve(backend, config):
+        """
+        Attempt to resolve a custom cloud backend.
+
+        Args:
+            backend: backend class
+            config: configuration parameters
+
+        Returns:
+            Cloud
+        """
+
+        try:
+            return Resolver()(backend)(config)
+
+        except Exception as e:
+            # Failure message
+            message = f'Unable to resolve cloud backend: "{backend}".'
+
+            # Append message if LIBCLOUD is not installed
+            message += ' Cloud storage is not available - install "cloud" extra to enable' if not LIBCLOUD else ""
+
+            raise ImportError(message) from e

txtai/cloud/hub.py

@@ -0,0 +1,101 @@
+"""
+Hugging Face Hub module
+"""
+
+import os
+import tempfile
+
+import huggingface_hub
+
+from huggingface_hub.utils import RepositoryNotFoundError
+
+from .base import Cloud
+
+
+class HuggingFaceHub(Cloud):
+    """
+    Hugging Face Hub cloud provider.
+    """
+
+    def metadata(self, path=None):
+        try:
+            # If this is an archive file path, get file metadata
+            if self.isarchive(path):
+                url = huggingface_hub.hf_hub_url(
+                    repo_id=self.config["container"], filename=os.path.basename(path), revision=self.config.get("revision")
+                )
+
+                return huggingface_hub.get_hf_file_metadata(url=url, token=self.config.get("token"))
+
+            # Otherwise return repository metadata
+            return huggingface_hub.model_info(repo_id=self.config["container"], revision=self.config.get("revision"), token=self.config.get("token"))
+
+        except RepositoryNotFoundError:
+            return None
+
+    def load(self, path=None):
+        # Download archvie file and return local path
+        if self.isarchive(path):
+            return huggingface_hub.hf_hub_download(
+                repo_id=self.config["container"],
+                filename=os.path.basename(path),
+                revision=self.config.get("revision"),
+                cache_dir=self.config.get("cache"),
+                token=self.config.get("token"),
+            )
+
+        # Download repository and return cached path
+        return huggingface_hub.snapshot_download(
+            repo_id=self.config["container"], revision=self.config.get("revision"), cache_dir=self.config.get("cache"), token=self.config.get("token")
+        )
+
+    def save(self, path):
+        # Get or create repository
+        huggingface_hub.create_repo(
+            repo_id=self.config["container"], token=self.config.get("token"), private=self.config.get("private", True), exist_ok=True
+        )
+
+        # Enable lfs-tracking of embeddings index files
+        self.lfstrack()
+
+        # Upload files
+        for f in self.listfiles(path):
+            huggingface_hub.upload_file(
+                repo_id=self.config["container"],
+                revision=self.config.get("revision"),
+                token=self.config.get("token"),
+                path_or_fileobj=f,
+                path_in_repo=os.path.basename(f),
+            )
+
+    def lfstrack(self):
+        """
+        Adds lfs-tracking of embeddings index files. This method adds tracking for documents and embeddings to .gitattributes.
+        """
+
+        # Get and read .gitattributes file
+        path = huggingface_hub.hf_hub_download(
+            repo_id=self.config["container"], filename=os.path.basename(".gitattributes"), token=self.config.get("token")
+        )
+
+        with open(path, "r", encoding="utf-8") as f:
+            content = f.read()
+
+        # Check if index files are lfs-tracked. Update .gitattributes, if necessary.
+        if "embeddings " not in content:
+            # Add documents and embeddings to lfs tracking
+            content += "documents filter=lfs diff=lfs merge=lfs -text\n"
+            content += "embeddings filter=lfs diff=lfs merge=lfs -text\n"
+
+            # pylint: disable=R1732
+            with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp:
+                tmp.write(content)
+                attributes = tmp.name
+
+            # Upload file
+            huggingface_hub.upload_file(
+                repo_id=self.config["container"], token=self.config.get("token"), path_or_fileobj=attributes, path_in_repo=os.path.basename(path)
+            )
+
+            # Remove temporary file
+            os.remove(attributes)

txtai/cloud/storage.py

@@ -0,0 +1,125 @@
+"""
+Object storage module
+"""
+
+import os
+
+# Conditional import
+try:
+    from libcloud.storage.providers import get_driver, DRIVERS
+    from libcloud.storage.types import ContainerDoesNotExistError, ObjectDoesNotExistError
+
+    LIBCLOUD = True
+except ImportError:
+    LIBCLOUD, DRIVERS = False, None
+
+
+from .base import Cloud
+
+
+class ObjectStorage(Cloud):
+    """
+    Object storage cloud provider backed by Apache libcloud.
+    """
+
+    @staticmethod
+    def isprovider(provider):
+        """
+        Checks if this provider is an object storage provider.
+
+        Args:
+            provider: provider name
+
+        Returns:
+            True if this is an object storage provider
+        """
+
+        return LIBCLOUD and provider and provider.lower() in [x.lower() for x in DRIVERS]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not LIBCLOUD:
+            raise ImportError('Cloud object storage is not available - install "cloud" extra to enable')
+
+        # Get driver for provider
+        driver = get_driver(config["provider"])
+
+        # Get client connection
+        self.client = driver(
+            config.get("key", os.environ.get("ACCESS_KEY")),
+            config.get("secret", os.environ.get("ACCESS_SECRET")),
+            **{field: config.get(field) for field in ["host", "port", "region", "token"] if config.get(field)},
+        )
+
+    def metadata(self, path=None):
+        try:
+            # If this is an archive path, check if file exists
+            if self.isarchive(path):
+                return self.client.get_object(self.config["container"], self.objectname(path))
+
+            # Otherwise check if container exists
+            return self.client.get_container(self.config["container"])
+        except (ContainerDoesNotExistError, ObjectDoesNotExistError):
+            return None
+
+    def load(self, path=None):
+        # Download archive file
+        if self.isarchive(path):
+            obj = self.client.get_object(self.config["container"], self.objectname(path))
+
+            # Create local directory, if necessary
+            directory = os.path.dirname(path)
+            if directory:
+                os.makedirs(directory, exist_ok=True)
+
+            obj.download(path, overwrite_existing=True)
+
+        # Download files in container. Optionally filter with a provided prefix.
+        else:
+            container = self.client.get_container(self.config["container"])
+            for obj in container.list_objects(prefix=self.config.get("prefix")):
+                # Derive local path and directory
+                localpath = os.path.join(path, obj.name)
+                directory = os.path.dirname(localpath)
+
+                # Create local directory, if necessary
+                os.makedirs(directory, exist_ok=True)
+
+                # Download file locally
+                obj.download(localpath, overwrite_existing=True)
+
+        return path
+
+    def save(self, path):
+        # Get or create container
+        try:
+            container = self.client.get_container(self.config["container"])
+        except ContainerDoesNotExistError:
+            container = self.client.create_container(self.config["container"])
+
+        # Upload files
+        for f in self.listfiles(path):
+            with open(f, "rb") as iterator:
+                self.client.upload_object_via_stream(iterator=iterator, container=container, object_name=self.objectname(f))
+
+    def objectname(self, name):
+        """
+        Derives an object name. This method checks if a prefix configuration parameter is present and combines
+        it with the input name parameter.
+
+        Args:
+            name: input name
+
+        Returns:
+            object name
+        """
+
+        # Get base name
+        name = os.path.basename(name)
+
+        # Get optional prefix/folder
+        prefix = self.config.get("prefix")
+
+        # Prepend prefix, if applicable
+        return f"{prefix}/{name}" if prefix else name

txtai/console/__init__.py

@@ -0,0 +1,5 @@
+"""
+Console imports
+"""
+
+from .base import Console