palimpzest 0.5.4__py3-none-any.whl → 0.6.1__py3-none-any.whl

palimpzest/core/data/datasources.py

@@ -1,369 +0,0 @@
-from __future__ import annotations
-
-import abc
-import base64
-import json
-import os
-import sys
-from io import BytesIO
-from typing import Any, Callable
-
-import modal
-import pandas as pd
-from bs4 import BeautifulSoup
-from papermage import Document
-
-from palimpzest import constants
-from palimpzest.core.elements.records import DataRecord
-from palimpzest.core.lib.schemas import (
-    DefaultSchema,
-    File,
-    ImageFile,
-    PDFFile,
-    Schema,
-    TextFile,
-    WebPage,
-    XLSFile,
-)
-from palimpzest.tools.pdfparser import get_text_from_pdf
-
-
-# First level of abstraction
-class AbstractDataSource(abc.ABC):
-    """
-    An AbstractDataSource is an Iterable which yields DataRecords adhering to a given schema.
-
-    This base class must have its `__iter__` method implemented by a subclass, with each
-    subclass reading data files from some real-world source (i.e. a directory, an S3 prefix,
-    etc.).
-
-    Many (if not all) DataSources should use Schemas from `palimpzest.elements.core`.
-    In the future, programmers can implement their own DataSources using custom Schemas.
-    """
-
-    def __init__(self, schema: Schema) -> None:
-        self._schema = schema
-
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}(schema={self.schema})"
-
-    def __eq__(self, __value: object) -> bool:
-        return self.__dict__ == __value.__dict__
-
-    @abc.abstractmethod
-    def __len__(self) -> int: ...
-
-    @abc.abstractmethod
-    def get_item(self, idx: int) -> DataRecord: ...
-
-    @abc.abstractmethod
-    def get_size(self) -> int: ...
-
-    @property
-    def schema(self) -> Schema:
-        return self._schema
-
-    def copy(self) -> AbstractDataSource:
-        raise NotImplementedError("You are calling this method from an abstract class.")
-
-    def serialize(self) -> dict[str, Any]:
-        return {"schema": self._schema.json_schema()}
-
-
-class DataSource(AbstractDataSource):
-    def __init__(self, schema: Schema, dataset_id: str) -> None:
-        super().__init__(schema)
-        self.dataset_id = dataset_id
-
-    def universal_identifier(self):
-        """Return a unique identifier for this Set."""
-        return self.dataset_id
-
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}(schema={self.schema}, dataset_id={self.dataset_id})"
-
-
-# Second level of abstraction
-class DirectorySource(DataSource):
-    """DirectorySource returns multiple File objects from a real-world source (a directory on disk)"""
-
-    def __init__(self, path: str, dataset_id: str, schema: Schema) -> None:
-        self.filepaths = [
-            os.path.join(path, filename)
-            for filename in sorted(os.listdir(path))
-            if os.path.isfile(os.path.join(path, filename))
-        ]
-        self.path = path
-        super().__init__(schema, dataset_id)
-
-    def serialize(self) -> dict[str, Any]:
-        return {
-            "schema": self.schema.json_schema(),
-            "path": self.path,
-            "source_type": "directory",
-        }
-
-    def __len__(self):
-        return len(self.filepaths)
-
-    def get_size(self):
-        # Get the memory size of the files in the directory
-        return sum([os.path.getsize(filepath) for filepath in self.filepaths])
-
-    def get_item(self, idx: int):
-        raise NotImplementedError("You are calling this method from an abstract class.")
-
-
-class FileSource(DataSource):
-    """FileSource returns a single File object from a single real-world local file"""
-
-    def __init__(self, path: str, dataset_id: str) -> None:
-        super().__init__(File, dataset_id)
-        self.filepath = path
-
-    def copy(self):
-        return FileSource(self.filepath, self.dataset_id)
-
-    def serialize(self) -> dict[str, Any]:
-        return {
-            "schema": self.schema.json_schema(),
-            "path": self.filepath,
-            "source_type": "file",
-        }
-
-    def __len__(self):
-        return 1
-
-    def get_size(self):
-        # Get the memory size of the filepath
-        return os.path.getsize(self.filepath)
-
-    def get_item(self, idx: int) -> DataRecord:
-        dr = DataRecord(self.schema, source_id=self.filepath)
-        dr.filename = self.filepath
-        with open(self.filepath, "rb") as f:
-            dr.contents = f.read()
-
-        return dr
-
-
-class MemorySource(DataSource):
-    """MemorySource returns multiple objects that reflect contents of an in-memory Python list
-        TODO(gerardo): Add support for other types of in-memory data structures (he has some code
-                   for subclassing MemorySource on his branch)
-    """
-
-    def __init__(self, vals: Any, dataset_id: str = "default_memory_input"):
-        if isinstance(vals, (str, int, float)):
-            self.vals = [vals]
-        elif isinstance(vals, tuple):
-            self.vals = list(vals)
-        else:
-            self.vals = vals
-        schema = Schema.from_df(self.vals) if isinstance(self.vals, pd.DataFrame) else DefaultSchema
-        super().__init__(schema, dataset_id)
-
-    def copy(self):
-        return MemorySource(self.vals, self.dataset_id)
-
-    def __len__(self):
-        return len(self.vals)
-
-    def get_size(self):
-        return sum([sys.getsizeof(self.get_item(idx)) for idx in range(len(self))])
-
-    def get_item(self, idx: int) -> DataRecord:
-        dr = DataRecord(self.schema, source_id=idx)
-        if isinstance(self.vals, pd.DataFrame):
-            row = self.vals.iloc[idx]
-            for field_name in row.index:
-                field_name_str = f"column_{field_name}" if isinstance(field_name, (int, float)) else str(field_name)
-                setattr(dr, field_name_str, row[field_name])
-        else:
-            dr.value = self.vals[idx]
-
-        return dr
-
-
-# Third level of abstraction
-class HTMLFileDirectorySource(DirectorySource):
-    def __init__(self, path: str, dataset_id: str) -> None:
-        super().__init__(path=path, dataset_id=dataset_id, schema=WebPage)
-        assert all([filename.endswith(tuple(constants.HTML_EXTENSIONS)) for filename in self.filepaths])
-
-    def copy(self):
-        return HTMLFileDirectorySource(self.path, self.dataset_id)
-
-    def html_to_text_with_links(self, html):
-        # Parse the HTML content
-        soup = BeautifulSoup(html, "html.parser")
-
-        # Find all hyperlink tags
-        for a in soup.find_all("a"):
-            # Check if the hyperlink tag has an 'href' attribute
-            if a.has_attr("href"):
-                # Replace the hyperlink with its text and URL in parentheses
-                a.replace_with(f"{a.text} ({a['href']})")
-
-        # Extract text from the modified HTML
-        text = soup.get_text(separator="\n", strip=True)
-        return text
-
-    def get_item(self, idx: int) -> DataRecord:
-        filepath = self.filepaths[idx]
-        dr = DataRecord(self.schema, source_id=filepath)
-        dr.filename = os.path.basename(filepath)
-        with open(filepath) as f:
-            text_content = f.read()
-
-        html = text_content
-        tokens = html.split()[: constants.MAX_HTML_ROWS]
-        dr.html = " ".join(tokens)
-
-        stripped_html = self.html_to_text_with_links(text_content)
-        tokens = stripped_html.split()[: constants.MAX_HTML_ROWS]
-        dr.text = " ".join(tokens)
-
-        return dr
-
-
-class ImageFileDirectorySource(DirectorySource):
-    def __init__(self, path: str, dataset_id: str) -> None:
-        super().__init__(path=path, dataset_id=dataset_id, schema=ImageFile)
-        assert all([filename.endswith(tuple(constants.IMAGE_EXTENSIONS)) for filename in self.filepaths])
-
-    def copy(self):
-        return ImageFileDirectorySource(self.path, self.dataset_id)
-
-    def get_item(self, idx: int) -> DataRecord:
-        filepath = self.filepaths[idx]
-        dr = DataRecord(self.schema, source_id=filepath)
-        dr.filename = os.path.basename(filepath)
-        with open(filepath, "rb") as f:
-            dr.contents = base64.b64encode(f.read())
-        return dr
-
-
-class PDFFileDirectorySource(DirectorySource):
-    def __init__(
-        self,
-        path: str,
-        dataset_id: str,
-        pdfprocessor: str = "modal",
-        file_cache_dir: str = "/tmp",
-    ) -> None:
-        super().__init__(path=path, dataset_id=dataset_id, schema=PDFFile)
-        assert all([filename.endswith(tuple(constants.PDF_EXTENSIONS)) for filename in self.filepaths])
-        self.pdfprocessor = pdfprocessor
-        self.file_cache_dir = file_cache_dir
-
-    def copy(self):
-        return PDFFileDirectorySource(self.path, self.dataset_id)
-
-    def get_item(self, idx: int) -> DataRecord:
-        filepath = self.filepaths[idx]
-        pdf_filename = os.path.basename(filepath)
-        with open(filepath, "rb") as f:
-            pdf_bytes = f.read()
-
-        if self.pdfprocessor == "modal":
-            print("handling PDF processing remotely")
-            remote_func = modal.Function.lookup("palimpzest.tools", "processPapermagePdf")
-        else:
-            remote_func = None
-
-        # generate text_content from PDF
-        if remote_func is not None:
-            doc_json_str = remote_func.remote([pdf_bytes])
-            docdict = json.loads(doc_json_str[0])
-            doc = Document.from_json(docdict)
-            text_content = ""
-            for p in doc.pages:
-                text_content += p.text
-        else:
-            text_content = get_text_from_pdf(pdf_filename, pdf_bytes, pdfprocessor=self.pdfprocessor, file_cache_dir=self.file_cache_dir)
-
-        # construct data record
-        dr = DataRecord(self.schema, source_id=filepath)
-        dr.filename = pdf_filename
-        dr.contents = pdf_bytes
-        dr.text_contents = text_content
-
-        return dr
-
-
-class TextFileDirectorySource(DirectorySource):
-    def __init__(self, path: str, dataset_id: str) -> None:
-        super().__init__(path=path, dataset_id=dataset_id, schema=TextFile)
-
-    def copy(self):
-        return TextFileDirectorySource(self.path, self.dataset_id)
-
-    def get_item(self, idx: int) -> DataRecord:
-        filepath = self.filepaths[idx]
-        dr = DataRecord(self.schema, source_id=filepath)
-        dr.filename = os.path.basename(filepath)
-        with open(filepath) as f:
-            dr.contents = f.read()
-        return dr
-
-
-class XLSFileDirectorySource(DirectorySource):
-    def __init__(self, path: str, dataset_id: str) -> None:
-        super().__init__(path=path, dataset_id=dataset_id, schema=XLSFile)
-        assert all([filename.endswith(tuple(constants.XLS_EXTENSIONS)) for filename in self.filepaths])
-
-    def copy(self):
-        return XLSFileDirectorySource(self.path, self.dataset_id)
-
-    def get_item(self, idx: int) -> DataRecord:
-        filepath = self.filepaths[idx]
-        dr = DataRecord(self.schema, source_id=filepath)
-        dr.filename = os.path.basename(filepath)
-        with open(filepath, "rb") as f:
-            dr.contents = f.read()
-
-        xls = pd.ExcelFile(BytesIO(dr.contents), engine="openpyxl")
-        dr.number_sheets = len(xls.sheet_names)
-        dr.sheet_names = xls.sheet_names
-        return dr
-
-
-# User-defined datasources
-class UserSource(DataSource):
-    """UserSource is a DataSource that is created by the user and not loaded from a file"""
-
-    def __init__(self, schema: Schema, dataset_id: str) -> None:
-        super().__init__(schema, dataset_id)
-
-    def serialize(self) -> dict[str, Any]:
-        return {
-            "schema": self.schema.json_schema(),
-            "source_type": "user-defined:" + self.__class__.__name__,
-        }
-
-    def __len__(self):
-        raise NotImplementedError("User needs to implement this method")
-
-    def get_size(self):
-        raise NotImplementedError("User may optionally implement this method.")
-
-    def get_item(self, idx: int) -> DataRecord:
-        raise NotImplementedError("User needs to implement this method.")
-
-    def copy(self):
-        raise NotImplementedError("User needs to implement this method.")
-
-class ValidationDataSource(UserSource):
-    """
-    TODO: update this class interface (and comment)
-    """
-
-    def get_val_length(self) -> int:
-        raise NotImplementedError("User needs to implement this method.")
-
-    def get_field_to_metric_fn(self) -> Callable:
-        raise NotImplementedError("User needs to implement this method.")
-
-    def get_item(self, idx: int, val: bool = False, include_label: bool = False) -> DataRecord:
-        raise NotImplementedError("User needs to implement this method.")

palimpzest/datamanager/datamanager.py

@@ -1,300 +0,0 @@
-import os
-import pickle
-from threading import Lock
-
-import pandas as pd
-import yaml
-
-from palimpzest import constants
-from palimpzest.config import Config
-from palimpzest.constants import DEFAULT_DATASET_ID_CHARS, MAX_DATASET_ID_CHARS, PZ_DIR
-from palimpzest.core.data.datasources import (
-    DataSource,
-    FileSource,
-    HTMLFileDirectorySource,
-    ImageFileDirectorySource,
-    MemorySource,
-    PDFFileDirectorySource,
-    TextFileDirectorySource,
-    UserSource,
-    XLSFileDirectorySource,
-)
-from palimpzest.utils.hash_helpers import hash_for_id
-
-
-class DataDirectorySingletonMeta(type):
-    _instances = {}
-    _lock: Lock = Lock()
-
-    def __call__(cls, *args, **kwargs):
-        with cls._lock:
-            if cls not in cls._instances:
-                instance = super().__call__(*args, **kwargs)
-                cls._instances[cls] = instance
-        return cls._instances[cls]
-
-
-class CacheService:
-    """This class manages the cache for the DataDirectory and other misc PZ components.
-    Eventually modify this to be durable and to have expiration policies."""
-
-    def __init__(self):
-        self.all_caches = {}
-
-    def get_cached_data(self, cache_name, cache_key):
-        return self.all_caches.setdefault(cache_name, {}).get(cache_key, None)
-
-    def put_cached_data(self, cache_name, cache_key, cache_val):
-        self.all_caches.setdefault(cache_name, {})[cache_key] = cache_val
-
-    def rm_cached_data(self, cache_name):
-        if cache_name in self.all_caches:
-            del self.all_caches[cache_name]
-
-    def rm_cache(self):
-        self.all_caches = {}
-
-
-class DataDirectory(metaclass=DataDirectorySingletonMeta):
-    """The DataDirectory is a registry of data sources."""
-
-    def __init__(self):
-        self._registry = {}
-        self._tempRegistry = {}
-        self._cache = {}
-        self._tempCache = {}
-        self.cacheService = CacheService()
-
-        # set up data directory
-        self._dir = PZ_DIR
-        current_config_path = os.path.join(self._dir, "current_config.yaml")
-        if not os.path.exists(self._dir):
-            os.makedirs(self._dir)
-            os.makedirs(self._dir + "/data/registered")
-            os.makedirs(self._dir + "/data/cache")
-            with open(self._dir + "/data/cache/registry.pkl", "wb") as f:
-                pickle.dump(self._registry, f)
-
-            # create default config
-            default_config = Config("default")
-            default_config.set_current_config()
-
-        # read current config (and dict. of configs) from disk
-        self._current_config = None
-        if os.path.exists(current_config_path):
-            with open(current_config_path) as f:
-                current_config_dict = yaml.safe_load(f)
-                self._current_config = Config(current_config_dict["current_config_name"])
-
-        # initialize the file cache directory, defaulting to the system's temporary directory "tmp/pz"
-        pz_file_cache_dir = self.current_config.get("filecachedir")
-        if pz_file_cache_dir and not os.path.exists(pz_file_cache_dir):
-            os.makedirs(pz_file_cache_dir)
-
-        # Unpickle the registry of data sources
-        if os.path.exists(self._dir + "/data/cache/registry.pkl"):
-            with open(self._dir + "/data/cache/registry.pkl", "rb") as f:
-                self._registry = pickle.load(f)
-
-        # Iterate through all items in the cache directory, and rebuild the table of entries
-        for root, _, files in os.walk(self._dir + "/data/cache"):
-            for file in files:
-                if file.endswith(".cached"):
-                    cache_id = file[:-7]
-                    self._cache[cache_id] = root + "/" + file
-
-    @property
-    def current_config(self):
-        if not self._current_config:
-            raise Exception("No current config found.")
-        return self._current_config
-
-    def get_cache_service(self):
-        return self.cacheService
-
-    def get_config(self):
-        return self.current_config._load_config()
-
-    def get_file_cache_dir(self):
-        return self.current_config.get("filecachedir")
-
-    #
-    # These methods handle properly registered data files, meant to be kept over the long haul
-    #
-    def register_local_directory(self, path, dataset_id):
-        """Register a local directory as a data source."""
-        self._registry[dataset_id] = ("dir", path)
-        with open(self._dir + "/data/cache/registry.pkl", "wb") as f:
-            pickle.dump(self._registry, f)
-
-    def register_local_file(self, path, dataset_id):
-        """Register a local file as a data source."""
-        self._registry[dataset_id] = ("file", path)
-        with open(self._dir + "/data/cache/registry.pkl", "wb") as f:
-            pickle.dump(self._registry, f)
-
-    def get_or_register_local_source(self, dataset_id_or_path):
-        """Return a dataset from the registry."""
-        if dataset_id_or_path in self._tempRegistry or dataset_id_or_path in self._registry:
-            return self.get_registered_dataset(dataset_id_or_path)
-        else:
-            if os.path.isfile(dataset_id_or_path):
-                self.register_local_file(dataset_id_or_path, dataset_id_or_path)
-            elif os.path.isdir(dataset_id_or_path):
-                self.register_local_directory(dataset_id_or_path, dataset_id_or_path)
-            else:
-                raise Exception(f"Path {dataset_id_or_path} is invalid. Does not point to a file or directory.")
-            return self.get_registered_dataset(dataset_id_or_path)
-
-    #TODO: need to revisit how to best leverage cache for memory sources.
-    def register_memory_source(self, vals, dataset_id):
-        """Register an in-memory dataset as a data source"""
-        self._tempRegistry[dataset_id] = ("memory", vals)
-
-    def get_or_register_memory_source(self, vals):
-        dataset_id = hash_for_id(str(vals), max_chars=DEFAULT_DATASET_ID_CHARS)
-        if dataset_id in self._tempRegistry:
-            return self.get_registered_dataset(dataset_id)
-        else:
-            self.register_memory_source(vals, dataset_id)
-        return self.get_registered_dataset(dataset_id)
-
-    def register_user_source(self, src: UserSource, dataset_id: str):
-        """Register a user source as a data source."""
-        # user sources are always ephemeral
-        self._tempRegistry[dataset_id] = ("user", src)
-
-    def get_registered_dataset(self, dataset_id):
-        """Return a dataset from the registry."""
-        if dataset_id in self._tempRegistry:
-            entry, rock = self._tempRegistry[dataset_id]
-        elif dataset_id in self._registry:
-            entry, rock = self._registry[dataset_id]
-        else:
-            raise Exception(f"Dataset {dataset_id} not found in the registry.")
-
-        if entry == "dir":
-            if all([f.endswith(tuple(constants.IMAGE_EXTENSIONS)) for f in os.listdir(rock)]):
-                return ImageFileDirectorySource(rock, dataset_id)
-            elif all([f.endswith(tuple(constants.PDF_EXTENSIONS)) for f in os.listdir(rock)]):
-                pdfprocessor = self.current_config.get("pdfprocessor")
-                if not pdfprocessor:
-                    raise Exception("No PDF processor found in the current config.")
-                file_cache_dir = self.get_file_cache_dir()
-                if not file_cache_dir:
-                    raise Exception("No file cache directory found.")
-                return PDFFileDirectorySource(
-                    path=rock, dataset_id=dataset_id, pdfprocessor=pdfprocessor, file_cache_dir=file_cache_dir
-                )
-            elif all([f.endswith(tuple(constants.XLS_EXTENSIONS)) for f in os.listdir(rock)]):
-                return XLSFileDirectorySource(rock, dataset_id)
-            elif all([f.endswith(tuple(constants.HTML_EXTENSIONS)) for f in os.listdir(rock)]):
-                return HTMLFileDirectorySource(rock, dataset_id)
-            else:
-                return TextFileDirectorySource(rock, dataset_id)
-
-        elif entry == "file":
-            return FileSource(rock, dataset_id)
-        elif entry == "memory":
-            return MemorySource(rock, dataset_id)
-        elif entry == "user":
-            src = rock
-            return src
-        else:
-            raise Exception("Unknown entry type")
-
-    def get_registered_dataset_type(self, dataset_id):
-        """Return the type of the given dataset in the registry."""
-        if dataset_id in self._tempRegistry:
-            entry, _ = self._tempRegistry[dataset_id]
-        elif dataset_id in self._registry:
-            entry, _ = self._registry[dataset_id]
-        else:
-            raise Exception("Cannot find dataset", dataset_id, "in the registry.")
-
-        return entry
-
-    def list_registered_datasets(self):
-        """Return a list of registered datasets."""
-        return self._registry.items()
-
-    def rm_registered_dataset(self, dataset_id):
-        """Remove a dataset from the registry."""
-        del self._registry[dataset_id]
-        with open(self._dir + "/data/cache/registry.pkl", "wb") as f:
-            pickle.dump(self._registry, f)
-
-    #
-    # These methods handle cached results. They are meant to be persisted for performance reasons,
-    # but can always be recomputed if necessary.
-    #
-    def get_cached_result(self, cache_id):
-        """Return a cached result."""
-        cached_result = None
-        if cache_id not in self._cache:
-            return cached_result
-
-        with open(self._cache[cache_id], "rb") as f:
-            cached_result = pickle.load(f)
-
-        return MemorySource(cached_result, cache_id)
-
-    def clear_cache(self, keep_registry=False):
-        """Clear the cache."""
-        self._cache = {}
-        self._tempCache = {}
-
-        # Delete all files in the cache directory (except registry.pkl if keep_registry=True)
-        for root, _, files in os.walk(self._dir + "/data/cache"):
-            for file in files:
-                if os.path.basename(file) != "registry.pkl" or keep_registry is False:
-                    os.remove(root + "/" + file)
-
-    def has_cached_answer(self, cache_id):
-        """Check if a dataset is in the cache."""
-        return cache_id in self._cache
-
-    def open_cache(self, cache_id):
-        if cache_id is not None and cache_id not in self._cache and cache_id not in self._tempCache:
-            self._tempCache[cache_id] = []
-            return True
-        return False
-
-    def append_cache(self, cache_id, data):
-        self._tempCache[cache_id].append(data)
-
-    def close_cache(self, cache_id):
-        """Close the cache."""
-        filename = self._dir + "/data/cache/" + cache_id + ".cached"
-        try:
-            with open(filename, "wb") as f:
-                pickle.dump(self._tempCache[cache_id], f)
-        except pickle.PicklingError:
-            print("Warning: Failed to save cache due to pickling error")
-            os.remove(filename)
-        del self._tempCache[cache_id]
-        self._cache[cache_id] = filename
-
-    def exists(self, dataset_id):
-        return dataset_id in self._registry
-
-    def get_path(self, dataset_id):
-        if dataset_id not in self._registry:
-            raise Exception("Cannot find dataset", dataset_id, "in the registry.")
-        entry, path = self._registry[dataset_id]
-        return path
-
-    def get_or_register_dataset(self, source: str | list | pd.DataFrame | DataSource):
-        if isinstance(source, str):
-            if len(source) > MAX_DATASET_ID_CHARS:
-                raise Exception(f"""Dataset ID {source} is too long. Maximum length is {MAX_DATASET_ID_CHARS} characters. 
-                                If you're passing a string data source, please wrap it in a list or pd.DataFrame.""")
-            source = self.get_or_register_local_source(source)
-        elif isinstance(source, (list, pd.DataFrame)):
-            source = self.get_or_register_memory_source(source)
-        elif isinstance(source, DataSource):
-            pass
-        else:
-            raise Exception(f"Invalid source type: {type(source)}, We only support pd.DataFrame, list, and str")
-
-        return source