aspect-data 0.0.1__py3-none-any.whl

aspect/io.py

@@ -0,0 +1,189 @@
+
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, Optional, Union
+from functools import partial
+import hashlib
+import os
+import tempfile
+
+from carabiner import print_err
+
+from numpy.typing import ArrayLike
+
+if TYPE_CHECKING:
+    from datasets import Dataset, IterableDataset
+    from pandas import DataFrame
+else:
+    Dataset, DataFrame, IterableDataset = Any, Any, Any
+
+from .package_data import DEFAULT_CACHE
+
+
+DATASETS_PREFIX: str = "hf://datasets/"
+
+def hasher(s: str, n: int = 16) -> str:
+    if isinstance(s, str):
+        s = s.encode()
+    return hashlib.sha256(s).hexdigest()[:n]
+
+def _lock_path(cache_dir: str, key: str) -> str:
+    locks_dir = os.path.join(cache_dir, ".locks")
+    os.makedirs(locks_dir, exist_ok=True)
+    h = hasher(key.encode("utf-8"))
+    return os.path.join(locks_dir, f"{h}.lock")
+
+
+def _load_from_file(filename: str, cache: Optional[str] = None) -> Dataset:
+    from datasets import load_dataset, Dataset, DatasetDict
+    from filelock import FileLock
+
+    cache = cache or DEFAULT_CACHE
+    if filename.removesuffix(".gz").endswith((".csv", ".tsv", ".txt")):
+        read_f = partial(
+            load_dataset,
+            path="csv",
+            data_files=filename,
+            cache_dir=cache,
+            sep="," if filename.endswith((".csv", ".csv.gz")) else "\t",
+        )
+        lock_key = ("csv", read_f.keywords.get('sep'))
+    elif filename.endswith((".arrow", ".hd5", ".json", ".parquet", ".xml")):
+        _, ext = os.path.splitext(filename)
+        protocol = ext.lstrip(".")
+        read_f = partial(
+            load_dataset,
+            path=protocol,
+            data_files=filename,
+            cache_dir=cache,
+        )
+        lock_key = (protocol, "")
+    elif filename.endswith(".hf"):
+        read_f = partial(
+            Dataset.load_from_disk, 
+            dataset_path=filename,
+        )
+        lock_key = ("hf", "")
+    else:
+        raise IOError(f"Could not infer how to open '{filename}' from its extension.")
+
+    # If no cache, nothing sensible to lock on
+    if cache is None:
+        ds = read_f()
+
+    # Cross-task lock on the shared filesystem
+    lockfile = _lock_path(cache, "_".join(lock_key))
+    with FileLock(lockfile, timeout=60. * 60.):
+        ds = read_f()
+    if isinstance(ds, DatasetDict):
+        return ds["train"]
+    else:
+        return ds
+
+
+def _load_from_dataframe(
+    dataframe: Union[DataFrame, Mapping[str, ArrayLike], Iterable[Mapping[str, ArrayLike]]],
+    cache: Optional[str] = None
+) -> Dataset:
+    from pandas import DataFrame
+
+    if cache is None:
+        cache = DEFAULT_CACHE
+        print_err(f"Defaulting to cache: {cache}")
+    if not isinstance(dataframe, DataFrame):
+        dataframe = DataFrame(dataframe)
+
+    hash_name = hasher(dataframe.to_string())
+    with tempfile.TemporaryDirectory() as tmpdir:
+        csv_filename = os.path.join(tmpdir, f"{hash_name}.csv.gz")
+        dataframe.to_csv(csv_filename, index=False)
+        ds = _load_from_file(
+            csv_filename, 
+            cache=cache,
+        )
+    return ds
+
+
+def _get_ref_chunk(
+    s, 
+    sep: Optional[str] = None, 
+    all_seps: str = "@~:"
+) -> str:
+    if sep is not None:
+        if sep in s:
+            s = s.rpartition(sep)[-1]
+        else:
+            return None
+    for _sep in all_seps:
+        s = s.partition(_sep)[0]
+    return s
+
+
+def _resolve_hf_hub_dataset(
+    ref: str, 
+    cache: Optional[str] = None
+) -> Dataset:
+    from datasets import concatenate_datasets, load_dataset, DatasetDict
+
+    ref = ref.removeprefix(DATASETS_PREFIX).removeprefix("hf://")
+    seps = "@~:"
+    ver = _get_ref_chunk(ref, "@", all_seps=seps)
+    split = _get_ref_chunk(ref, ":", all_seps=seps)
+    config = _get_ref_chunk(ref, "~", all_seps=seps)
+    
+    ds = load_dataset(
+        path=_get_ref_chunk(ref, all_seps=seps), 
+        name=config, 
+        split=split, 
+        revision=ver, 
+        cache_dir=cache,
+    )
+    if isinstance(ds, DatasetDict):
+        ds = concatenate_datasets([v for key, v in ds.items()])
+    return ds
+
+
+class AutoDataset:
+
+    def __init__(self, dataset):
+        self._dataset = dataset
+
+    @classmethod
+    def load(
+        cls, 
+        data: Union[str, DataFrame], 
+        cache: Optional[str] = None
+    ) -> Union[Dataset, IterableDataset]:
+        from datasets import load_dataset, Dataset, IterableDataset
+        from pandas import DataFrame
+
+        if isinstance(data, (Dataset, IterableDataset)):
+            dataset = data
+        elif isinstance(data, (DataFrame, Mapping)):
+            dataset = _load_from_dataframe(
+                data, 
+                cache=cache,
+            )
+        elif isinstance(data, str):
+            if data.startswith("hf://"):
+                dataset = _resolve_hf_hub_dataset(
+                    data,
+                    cache=cache,
+                )
+            elif os.path.exists(data):
+                dataset = _load_from_file(
+                    data,
+                    cache=cache,
+                )
+            else:
+                raise ValueError(
+                    f"""
+                    If `data` is a string, it must start with "{DATASETS_PREFIX}" or a path to an existing file. 
+                    It was "{data}".
+                    """
+                )
+        else:
+            raise ValueError(
+                """
+                Data must be a string, Dataset, dictionary, or Pandas DataFrame.
+                """
+            )
+        return cls(dataset)

aspect/package_data.py

@@ -0,0 +1,35 @@
+"""Tools for loading and writing reusable package data."""
+
+import os
+
+from platformdirs import user_cache_dir
+
+from . import app_name, __author__, __version__
+
+CACHE_DIR = user_cache_dir(
+    app_name,
+    appauthor=False,
+    version=__version__,
+)
+
+ELUENT_CACHE = "ASPECT_CACHE"
+DEFAULT_CACHE = CACHE_DIR
+
+def _get_data_path(
+    filename: str, 
+    env_key: str = ELUENT_CACHE,
+    default: str = DEFAULT_CACHE
+) -> str:
+    """Returns the path to a writable version of a package data file.
+    
+    Copies it from the package resources if not present.
+
+    """
+    cache_dir = os.environ.get(
+        env_key, 
+        os.path.expanduser(default),
+    )
+    os.makedirs(cache_dir, exist_ok=True)
+    os.environ["HF_HOME"] = cache_dir
+    os.environ["HF_DATASETS_CACHE"] = cache_dir
+    return cache_dir, os.path.join(cache_dir, filename)

aspect/transform/__init__.py

@@ -0,0 +1,9 @@
+from .registry import FUNCTION_REGISTRY
+
+def _load_all():
+    try:
+        from . import functions, deep_functions  # importing populates FUNCTION_REGISTRY as a side effect
+    except (ImportError, NameError):
+        pass
+
+_load_all()

aspect/transform/base.py

@@ -0,0 +1,182 @@
+"""Classes to enable JSON serialization of preprocessing functions."""
+
+from typing import Any, Callable, Dict, Iterable, Mapping, Union
+from dataclasses import asdict, dataclass, field
+from functools import partial
+
+from carabiner import print_err
+import numpy as np
+
+from .. import app_name, __version__
+from ..checkpoint_utils import save_json, _load_json
+from .registry import FUNCTION_REGISTRY
+
+
+@dataclass
+class ColumnTransform:
+
+    """Serializable wrapper around a registered transform function.
+
+    Examples
+    ========
+    Basic construction and output column naming:
+
+    >>> ct = ColumnTransform(name="log", input_column="affinity")
+    >>> ct.input_column
+    'affinity'
+    >>> ct.name
+    'log'
+    >>> ct.output_column.startswith("aspect-data/")
+    True
+
+    Output column is deterministic:
+
+    >>> ct2 = ColumnTransform(name="log", input_column="affinity")
+    >>> ct.output_column == ct2.output_column
+    True
+
+    Output column changes with input_column:
+
+    >>> ct3 = ColumnTransform(name="log", input_column="other")
+    >>> ct.output_column != ct3.output_column
+    True
+
+    Output column changes with kwargs:
+
+    >>> ct4 = ColumnTransform(name="hash", input_column="assay", kwargs={"ndim": 64})
+    >>> ct5 = ColumnTransform(name="hash", input_column="assay", kwargs={"ndim": 128})
+    >>> ct4.output_column != ct5.output_column
+    True
+
+    Applying a transform:
+
+    >>> import numpy as np
+    >>> data = {"affinity": [1.0, 10.0, 100.0]}
+    >>> result = ct(data)
+    >>> ct.output_column in result
+    True
+    >>> np.allclose(result[ct.output_column], np.log([1., 10., 100.])[:,None])
+    True
+
+    Idempotency — calling twice skips recomputation:
+
+    >>> result2 = ct(result)
+    >>> np.array_equal(result[ct.output_column], result2[ct.output_column])
+    True
+
+    Serialization roundtrip:
+
+    >>> d = ct.to_dict()
+    >>> ct_restored = ColumnTransform.from_dict(d)
+    >>> ct_restored.output_column == ct.output_column
+    True
+    >>> data2 = {"affinity": [1.0, 10.0, 100.0]}
+    >>> result3 = ct_restored(data2)
+    >>> np.allclose(result3[ct.output_column], np.log([1., 10., 100.])[:,None])
+    True
+
+    Chained application propagates columns correctly:
+
+    >>> ct_hash = ColumnTransform(name="hash", input_column="assay", kwargs={"ndim": 16})
+    >>> data3 = {"assay": ["MIC", "MIC", "binding"]}
+    >>> r1 = ct_hash(data3)
+    >>> ct_hash.output_column in r1
+    True
+    >>> r1[ct_hash.output_column].shape
+    (3, 16)
+
+    Unregistered name raises:
+
+    >>> ColumnTransform(name="nonexistent", input_column="x")
+    Traceback (most recent call last):
+        ...
+    ValueError: Function 'nonexistent' is not registered...
+
+    show() lists registered functions:
+
+    >>> transforms = ColumnTransform.show()
+    >>> "log" in transforms
+    True
+    >>> "hash" in transforms
+    True
+
+    """
+
+    name: Union[str, Callable]
+    input_column: str
+    kwargs: Mapping[str, Any] = field(default_factory = dict)
+    _version: str = field(default=__version__)
+    _app: str = field(default=app_name)
+
+    def __post_init__(self):
+        if isinstance(self.name, str):
+            try:
+                self.closure = FUNCTION_REGISTRY[self.name]
+            except KeyError:
+                raise ValueError(
+                    f"Function '{self.name}' is not registered. Try: {', '.join(FUNCTION_REGISTRY)}"
+                )
+            self._serializable = True
+        elif isinstance(self.name, Callable):
+            self.closure = self.name
+            self.name = str(self.name)
+            self._serializable = False
+        else:
+            raise ValueError(f"Transformation function must be a str or callable")
+        self.hash = self._get_hash()
+        self.hash_stub = self.hash[:8]
+        self.output_column = self._get_output_column()
+        self.function = self.closure(**self.kwargs)
+
+    def _get_hash(self):
+        from datasets.fingerprint import Hasher
+        return Hasher.hash(self.to_dict())
+
+    def _get_output_column(self):
+        return f"{self._app}/{self._version}/{self.name}:{self.input_column}/{self.hash_stub}"
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the function call to a JSON-serializable format."""
+        if self._serializable:
+            return asdict(self)
+        else:
+            raise AttributeError(f"Cannot serialize {self} with callable {self.name}.")
+
+    def to_file(self, filename: str) -> None:
+        """Save to JSON."""
+        return save_json(self.to_dict(), filename)
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Union[str, Mapping]]) -> 'ColumnTransform':
+        """Reconstruct the function call from JSON."""
+        return cls(**data)
+
+    @classmethod
+    def from_file(cls, filename: str) -> 'ColumnTransform':
+        """Load from JSON."""
+        return cls.from_dict(_load_json(filename))
+
+    @classmethod
+    def show(cls) -> 'ColumnTransform':
+        """List registered functions."""
+        return tuple(FUNCTION_REGISTRY)
+
+    def apply(
+        self, 
+        inputs: Mapping[str, Iterable]
+    ) -> np.ndarray:
+        result = self.function(inputs, self.input_column)
+        if result.ndim == 1:
+            result = result[:, None]
+        return result
+
+    def __call__(
+        self, 
+        inputs: Mapping[str, Iterable]
+    ) -> Dict[str, np.ndarray]:
+        """Execute the function with stored parameters."""
+        if self.output_column not in inputs:
+            inputs[self.output_column] = self.apply(inputs)
+        else:
+            print_err(f"[WARN] {self.output_column} already present, skipping.")
+        return inputs

aspect/transform/deep_functions.py

@@ -0,0 +1,143 @@
+"""Preprocessing functions from deep embeddings."""
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Mapping, Tuple
+from functools import partial
+
+from carabiner import cast, print_err
+import numpy as np
+from numpy.typing import ArrayLike
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase, PreTrainedModel
+else:
+    PreTrainedTokenizerBase, PreTrainedModel = Any, Any
+
+from .registry import register_function
+from ..typing import StrOrIterableOfStr
+
+def get_device():
+    import torch
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _index_into(x: ArrayLike, i: int):
+    return x[:,i,:]
+
+
+def _get_value_from_tuple(f: Callable, **kwargs):
+    partial_f = partial(f, **kwargs)
+
+    def _f(*args, **fkwargs):
+        return partial_f(*args, **fkwargs).values
+
+    return _f
+
+
+def _resolve_bart_model(ref: str) -> Tuple[PreTrainedTokenizerBase, PreTrainedModel]:
+    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+    import torch
+    model = AutoModelForSeq2SeqLM.from_pretrained(ref)
+    return AutoTokenizer.from_pretrained(ref), torch.compile(model, fullgraph=True).to(get_device())
+
+
+def _tokenize_for_embedding(
+    x: Iterable[str],
+    tokenizer: PreTrainedTokenizerBase
+) -> Dict[str, np.ndarray]:
+
+    inputs = [
+        _x if _x is not None 
+        else '<unk>' 
+        for _x in cast(x, to=list)
+    ]
+    tokenizer_args = {
+        "return_tensors": "pt",
+        "padding": True,
+    }
+    try:
+        tokenized = tokenizer(inputs, **tokenizer_args)
+    except (ValueError, TypeError) as e:
+        print_err(e)
+        print_err(f"Tokenizing this batch failed: {inputs}")
+        tokenized = tokenizer(
+            ['<unk>'] * len(inputs), 
+            **tokenizer_args,
+        )
+
+    return {
+        key: tokenized[key].to(get_device()) 
+        for key in ['input_ids', 'attention_mask']
+    }
+
+
+@register_function("hf-bart")
+def HfBART(
+    ref: str,
+    aggregator: StrOrIterableOfStr = "mean"
+) -> np.ndarray:
+
+    try:
+        from transformers import PreTrainedTokenizerBase, PreTrainedModel
+    except ImportError:
+        raise ImportError("Hugging Face transformers not installed. Try `pip install aspect[deep]`.")
+    else:
+        import torch
+
+    EMBEDDING_AGGREGATORS = {
+        "start": partial(_index_into, i=0),
+        "end": partial(_index_into, i=-1),
+        "sum": partial(torch.sum, dim=-2),
+        "mean": partial(torch.mean, dim=-2),
+        "median": _get_value_from_tuple(torch.median, dim=-2),
+        "max": _get_value_from_tuple(torch.max, dim=-2),
+    }
+    DEFAULT_AGGREGATOR = EMBEDDING_AGGREGATORS["mean"]
+
+    def _aggregated_embedding(
+        x: torch.Tensor, 
+        aggregators: Iterable[Callable] = (DEFAULT_AGGREGATOR, )
+    ):
+        aggregated = [
+            a(x).detach().cpu().numpy() for a in aggregators
+        ]
+
+        return np.concatenate(aggregated, axis=-1)
+
+
+    aggregators = [a.casefold() for a in cast(aggregator, to=list)]
+    incorrect_agg = [
+        a for a in aggregators if a not in EMBEDDING_AGGREGATORS
+    ]
+    if len(incorrect_agg) > 0:
+        raise ValueError(f"Requested aggregators do not exist: {incorrect_agg}")
+    else:
+        aggregators = [EMBEDDING_AGGREGATORS[a] for a in aggregators]
+    tokenizer, model = _resolve_bart_model(ref)
+
+    def _hf_bart(
+        data: Mapping[str, Iterable],
+        input_column: str
+    ) -> np.ndarray:
+        tokenized_inputs = _tokenize_for_embedding(
+            data[input_column],
+            tokenizer=tokenizer,
+        )
+        model.eval()
+        with torch.no_grad():
+            outputs = model(
+                **tokenized_inputs, 
+                decoder_input_ids=tokenized_inputs['input_ids'],
+                output_hidden_states=True,
+                return_dict=True,
+            )
+        enc_last = _aggregated_embedding(
+            outputs.encoder_hidden_states[-1],
+            aggregators=aggregators,
+        )
+        dec_last = _aggregated_embedding(
+            outputs.decoder_hidden_states[-1],
+            aggregators=aggregators,
+        )
+        return np.concatenate([enc_last, dec_last], axis=-1)
+
+    return _hf_bart