tuft 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

tuft/loss_fn/__init__.py

@@ -0,0 +1,48 @@
+from typing import Callable, Dict, Tuple
+
+from torch import Tensor
+from typing_extensions import TypeAlias
+
+from ..exceptions import (
+    LossFunctionInputShapeMismatchException,
+    LossFunctionMissingInputException,
+    LossFunctionNotFoundException,
+)
+
+
+LossFnType: TypeAlias = Callable[
+    [Dict[str, Tensor], Dict[str, float]], Tuple[Tensor, Dict[str, float]]
+]
+
+LOSS_FN = {
+    "cispo": "tuft.loss_fn.cispo.cispo_loss",
+    "cross_entropy": "tuft.loss_fn.cross_entropy.cross_entropy_loss",
+    "dro": "tuft.loss_fn.dro.dro_loss",
+    "importance_sampling": "tuft.loss_fn.importance_sampling.importance_sampling_loss",
+    "ppo": "tuft.loss_fn.ppo.ppo_loss",
+}
+
+
+def get_loss_fn(loss_fn_name: str) -> LossFnType:
+    """Retrieve the loss function by name."""
+    if loss_fn_name not in LOSS_FN:
+        raise LossFunctionNotFoundException(loss_fn_name)
+
+    module_path, func_name = LOSS_FN[loss_fn_name].rsplit(".", 1)
+    module = __import__(module_path, fromlist=[func_name])
+    return getattr(module, func_name)
+
+
+def _check_loss_fn_inputs(
+    loss_fn_inputs: Dict[str, Tensor], required_keys: Tuple[str, ...], check_shapes: bool = False
+) -> None:
+    """Check if all required keys are present in loss_fn_inputs and optionally
+    check if their shapes match."""
+    for key in required_keys:
+        if key not in loss_fn_inputs:
+            raise LossFunctionMissingInputException(key)
+
+    if check_shapes:
+        shapes = [loss_fn_inputs[key].shape for key in required_keys]
+        if not all(shape == shapes[0] for shape in shapes):
+            raise LossFunctionInputShapeMismatchException(shapes)

tuft/loss_fn/cispo.py

@@ -0,0 +1,40 @@
+from typing import Dict, Tuple
+
+import torch
+
+from . import _check_loss_fn_inputs
+
+
+def cispo_loss(
+    loss_fn_inputs: Dict[str, torch.Tensor], loss_fn_config: Dict[str, float]
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """Computes the Clipped Importance Sampling Policy Optimization (CISPO) loss.
+
+    Args:
+        loss_fn_inputs: A dictionary of tensors required for the loss function.
+            Expected keys: "target_logprobs", "logprobs", "advantages".
+        loss_fn_config: A dictionary of configuration parameters for the loss function.
+            Expected keys: "clip_low_threshold", "clip_high_threshold".
+
+    Returns:
+        A tuple containing the computed loss and a dictionary of metrics.
+    """
+    _check_loss_fn_inputs(
+        loss_fn_inputs, ("target_logprobs", "logprobs", "advantages"), check_shapes=True
+    )
+    target_logprobs = loss_fn_inputs["target_logprobs"]
+    sampling_logprobs = loss_fn_inputs["logprobs"]
+    advantages = loss_fn_inputs["advantages"]
+    clip_low_threshold = loss_fn_config.get("clip_low_threshold", 0.9)
+    clip_high_threshold = loss_fn_config.get("clip_high_threshold", 1.1)
+
+    # Compute probability ratio
+    prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
+    # Apply clipping
+    clipped_ratio = torch.clamp(prob_ratio, clip_low_threshold, clip_high_threshold)
+    # Compute CISPO objective (detach the clipped ratio)
+    cispo_objective = clipped_ratio.detach() * target_logprobs * advantages
+    # CISPO loss is negative of objective
+    loss = -cispo_objective.sum()
+
+    return loss, {"loss:sum": loss.item()}

tuft/loss_fn/cross_entropy.py

@@ -0,0 +1,26 @@
+from typing import Dict, Tuple
+
+import torch
+
+from . import _check_loss_fn_inputs
+
+
+def cross_entropy_loss(
+    loss_fn_inputs: Dict[str, torch.Tensor], loss_fn_config: Dict[str, float]
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """Computes the Cross Entropy loss.
+
+    Args:
+        loss_fn_inputs: A dictionary of tensors required for the loss function.
+            Expected keys: "target_logprobs", "weights".
+        loss_fn_config: A dictionary of configuration parameters for the loss function.
+            (No expected keys for this loss function.)
+    Returns:
+        A tuple containing the computed loss and a dictionary of metrics.
+    """
+    _check_loss_fn_inputs(loss_fn_inputs, ("target_logprobs", "weights"), check_shapes=True)
+    target_logprobs = loss_fn_inputs["target_logprobs"]
+    weights = loss_fn_inputs["weights"]
+
+    loss = -(target_logprobs * weights).sum()
+    return loss, {"loss:sum": loss.item()}

tuft/loss_fn/dro.py

@@ -0,0 +1,37 @@
+from typing import Dict, Tuple
+
+import torch
+
+from . import _check_loss_fn_inputs
+
+
+def dro_loss(
+    loss_fn_inputs: Dict[str, torch.Tensor], loss_fn_config: Dict[str, float]
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """Computes the Distributionally Robust Optimization (DRO) loss.
+
+    Args:
+        loss_fn_inputs: A dictionary of tensors required for the loss function.
+            Expected keys: "target_logprobs", "logprobs", "advantages".
+        loss_fn_config: A dictionary of configuration parameters for the loss function.
+            Expected keys: "beta".
+
+    Returns:
+        A tuple containing the computed loss and a dictionary of metrics.
+    """
+    _check_loss_fn_inputs(
+        loss_fn_inputs, ("target_logprobs", "logprobs", "advantages"), check_shapes=True
+    )
+    target_logprobs = loss_fn_inputs["target_logprobs"]
+    sampling_logprobs = loss_fn_inputs["logprobs"]
+    advantages = loss_fn_inputs["advantages"]
+    beta = loss_fn_config.get("beta", 0.01)
+
+    # Compute quadratic penalty term
+    quadratic_term = (target_logprobs - sampling_logprobs) ** 2
+    # Compute DRO objective
+    dro_objective = target_logprobs * advantages - 0.5 * beta * quadratic_term
+    # DRO loss is negative of objective
+    loss = -dro_objective.sum()
+
+    return loss, {"loss:sum": loss.item()}

tuft/loss_fn/importance_sampling.py

@@ -0,0 +1,33 @@
+from typing import Dict, Tuple
+
+import torch
+
+from . import _check_loss_fn_inputs
+
+
+def importance_sampling_loss(
+    loss_fn_inputs: Dict[str, torch.Tensor], loss_fn_config: Dict[str, float]
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """Computes the importance sampling loss.
+
+    Args:
+        loss_fn_inputs: A dictionary of tensors required for the loss function.
+            Expected keys: "target_logprobs", "logprobs", "advantages".
+        loss_fn_config: This parameter is unused.
+
+    Returns:
+        A tuple containing the computed loss and a dictionary of metrics.
+    """
+    _check_loss_fn_inputs(
+        loss_fn_inputs, ("target_logprobs", "logprobs", "advantages"), check_shapes=True
+    )
+    target_logprobs = loss_fn_inputs["target_logprobs"]
+    sampling_logprobs = loss_fn_inputs["logprobs"]
+    advantages = loss_fn_inputs["advantages"]
+
+    # Compute probability ratio
+    prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
+    # Compute importance-weighted loss
+    loss = -(prob_ratio * advantages).sum()
+
+    return loss, {"loss:sum": loss.item()}

tuft/loss_fn/ppo.py

@@ -0,0 +1,43 @@
+from typing import Dict, Tuple
+
+import torch
+
+from . import _check_loss_fn_inputs
+
+
+def ppo_loss(
+    loss_fn_inputs: Dict[str, torch.Tensor], loss_fn_config: Dict[str, float]
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """Computes the Proximal Policy Optimization (PPO) loss.
+
+    Args:
+        loss_fn_inputs: A dictionary of tensors required for the loss function.
+            Expected keys: "target_logprobs", "logprobs", "advantages".
+        loss_fn_config: A dictionary of configuration parameters for the loss function.
+            Expected keys: "clip_low_threshold", "clip_high_threshold".
+
+    Returns:
+        A tuple containing the computed loss and a dictionary of metrics.
+    """
+    _check_loss_fn_inputs(
+        loss_fn_inputs, ("target_logprobs", "logprobs", "advantages"), check_shapes=True
+    )
+    target_logprobs = loss_fn_inputs["target_logprobs"]
+    sampling_logprobs = loss_fn_inputs["logprobs"]
+    advantages = loss_fn_inputs["advantages"]
+    clip_low_threshold = loss_fn_config.get("clip_low_threshold", 0.9)
+    clip_high_threshold = loss_fn_config.get("clip_high_threshold", 1.1)
+
+    # Compute probability ratio
+    prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
+    # Apply clipping
+    clipped_ratio = torch.clamp(prob_ratio, clip_low_threshold, clip_high_threshold)
+    # Compute both objectives
+    unclipped_objective = prob_ratio * advantages
+    clipped_objective = clipped_ratio * advantages
+    # Take minimum (most conservative)
+    ppo_objective = torch.min(unclipped_objective, clipped_objective)
+    # PPO loss is negative of objective
+    loss = -ppo_objective.sum()
+
+    return loss, {"loss:sum": loss.item()}

tuft/persistence/__init__.py

@@ -0,0 +1,32 @@
+"""Persistence package exports."""
+
+from __future__ import annotations
+
+from .redis_store import (
+    DEFAULT_FUTURE_TTL_SECONDS,
+    PersistenceConfig,
+    PersistenceMode,
+    RedisPipeline,
+    RedisStore,
+    delete_record,
+    get_redis_store,
+    is_persistence_enabled,
+    load_record,
+    save_record,
+    save_records_atomic,
+)
+
+
+__all__ = [
+    "DEFAULT_FUTURE_TTL_SECONDS",
+    "PersistenceConfig",
+    "PersistenceMode",
+    "RedisPipeline",
+    "RedisStore",
+    "delete_record",
+    "get_redis_store",
+    "is_persistence_enabled",
+    "load_record",
+    "save_record",
+    "save_records_atomic",
+]

tuft/persistence/file_redis.py

@@ -0,0 +1,268 @@
+"""File-backed Redis-like store for small demos and tests.
+
+This module implements a minimal subset of redis-py behaviors with a JSON
+backing file. It is designed for low-volume usage where performance is not a
+concern. All write operations flush the full in-memory state to disk.
+
+Example:
+    from pathlib import Path
+
+    from tuft.persistence.file_redis import FileRedis
+
+    store = FileRedis(Path("~/.cache/tuft/file_redis.json").expanduser())
+    store.set("alpha", "1")
+    store.setex("beta", 5, "2")
+    assert store.get("alpha") == "1"
+    for key in store.scan_iter(match="a*"):
+        print(key)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import time
+from dataclasses import dataclass
+from fnmatch import fnmatch
+from pathlib import Path
+from typing import Iterable
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _FileRedisValue:
+    value: str
+    expires_at: float | None
+
+
+class FileRedis:
+    """Tiny file-backed Redis-like store for tests and demos.
+
+    Args:
+        file_path: Path to the JSON file used for persistence.
+
+    Example:
+        from pathlib import Path
+
+        store = FileRedis(Path("/tmp/file_redis.json"))
+        store.set("key", "value")
+        assert store.get("key") == "value"
+    """
+
+    def __init__(self, file_path: Path) -> None:
+        self._file_path = Path(file_path)
+        self._file_path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._data: dict[str, _FileRedisValue] = {}
+        self._load()
+
+    def _load(self) -> None:
+        """Load persisted data from disk into memory."""
+        if not self._file_path.exists():
+            return
+        try:
+            raw = json.loads(self._file_path.read_text(encoding="utf-8"))
+            for key, payload in raw.items():
+                if not isinstance(payload, dict):
+                    continue
+                self._data[key] = _FileRedisValue(
+                    value=str(payload.get("value", "")),
+                    expires_at=payload.get("expires_at"),
+                )
+        except (json.JSONDecodeError, OSError):
+            logger.exception("Failed to load FileRedis data from %s", self._file_path)
+
+    def _dump(self) -> None:
+        """Write the in-memory store to disk as JSON."""
+        payload = {
+            key: {"value": entry.value, "expires_at": entry.expires_at}
+            for key, entry in self._data.items()
+        }
+        tmp_path = self._file_path.with_suffix(self._file_path.suffix + ".tmp")
+        tmp_path.write_text(json.dumps(payload, ensure_ascii=True), encoding="utf-8")
+        tmp_path.replace(self._file_path)
+
+    def _purge_expired(self) -> None:
+        """Remove expired keys and persist the updated store."""
+        now = time.time()
+        expired = [
+            key for key, entry in self._data.items() if entry.expires_at and entry.expires_at <= now
+        ]
+        if expired:
+            for key in expired:
+                self._data.pop(key, None)
+            self._dump()
+
+    def set(self, key: str, value: str) -> bool:
+        """Set a key to a string value.
+
+        Args:
+            key: Key to set.
+            value: String value to store.
+
+        Returns:
+            True on success.
+        """
+        with self._lock:
+            self._data[key] = _FileRedisValue(value=value, expires_at=None)
+            self._dump()
+        return True
+
+    def setex(self, key: str, ttl_seconds: int | float, value: str) -> bool:
+        """Set a key with TTL in seconds.
+
+        Args:
+            key: Key to set.
+            ttl_seconds: Time-to-live in seconds.
+            value: String value to store.
+
+        Returns:
+            True on success.
+        """
+        with self._lock:
+            expires_at = time.time() + float(ttl_seconds)
+            self._data[key] = _FileRedisValue(value=value, expires_at=expires_at)
+            self._dump()
+        return True
+
+    def get(self, key: str) -> str | None:
+        """Get a string value by key.
+
+        Args:
+            key: Key to retrieve.
+
+        Returns:
+            The stored value, or None if missing/expired.
+        """
+        with self._lock:
+            self._purge_expired()
+            entry = self._data.get(key)
+            return entry.value if entry else None
+
+    def delete(self, *keys: str) -> int:
+        """Delete one or more keys.
+
+        Args:
+            *keys: Keys to delete.
+
+        Returns:
+            Number of keys removed.
+        """
+        removed = 0
+        with self._lock:
+            for key in keys:
+                if key in self._data:
+                    self._data.pop(key, None)
+                    removed += 1
+            if removed:
+                self._dump()
+        return removed
+
+    def exists(self, key: str) -> int:
+        """Check if a key exists.
+
+        Args:
+            key: Key to check.
+
+        Returns:
+            1 if the key exists, otherwise 0.
+        """
+        with self._lock:
+            self._purge_expired()
+            return 1 if key in self._data else 0
+
+    def scan_iter(self, match: str | None = None) -> Iterable[str]:
+        """Iterate over keys matching a pattern.
+
+        Args:
+            match: Optional glob pattern (e.g., "prefix:*").
+
+        Returns:
+            An iterator of matching keys.
+        """
+        with self._lock:
+            self._purge_expired()
+            keys = list(self._data.keys())
+        pattern = match or "*"
+        for key in keys:
+            if fnmatch(key, pattern):
+                yield key
+
+    def pipeline(self, transaction: bool = True) -> "FileRedisPipeline":
+        """Create a pipeline for batched operations.
+
+        Args:
+            transaction: Ignored; kept for compatibility.
+
+        Returns:
+            A FileRedisPipeline instance.
+        """
+        _ = transaction  # kept for signature compatibility
+        return FileRedisPipeline(self)
+
+    def close(self) -> None:
+        """No-op close for API compatibility."""
+        return None
+
+
+class FileRedisPipeline:
+    """Minimal pipeline that writes once on exit.
+
+    Example:
+        with store.pipeline() as pipe:
+            pipe.set("a", "1")
+            pipe.setex("b", 10, "2")
+            pipe.delete("c")
+    """
+
+    def __init__(self, store: FileRedis) -> None:
+        self._store = store
+        self._ops: list[tuple[str, tuple]] = []
+
+    def __enter__(self) -> "FileRedisPipeline":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if exc_type is None:
+            self._execute()
+
+    def set(self, key: str, value: str) -> "FileRedisPipeline":
+        """Queue a SET operation."""
+        self._ops.append(("set", (key, value)))
+        return self
+
+    def setex(self, key: str, ttl_seconds: int | float, value: str) -> "FileRedisPipeline":
+        """Queue a SETEX operation."""
+        self._ops.append(("setex", (key, ttl_seconds, value)))
+        return self
+
+    def delete(self, *keys: str) -> "FileRedisPipeline":
+        """Queue a DELETE operation."""
+        self._ops.append(("delete", keys))
+        return self
+
+    def _execute(self) -> None:
+        """Apply queued operations and flush to disk."""
+        with self._store._lock:
+            for op, args in self._ops:
+                if op == "set":
+                    key, value = args
+                    self._store._data[key] = _FileRedisValue(value=value, expires_at=None)
+                elif op == "setex":
+                    key, ttl_seconds, value = args
+                    expires_at = time.time() + float(ttl_seconds)
+                    self._store._data[key] = _FileRedisValue(value=value, expires_at=expires_at)
+                elif op == "delete":
+                    for key in args:
+                        self._store._data.pop(key, None)
+            if self._ops:
+                self._store._dump()
+            self._ops.clear()
+
+    def execute(self) -> list[object]:
+        """Execute queued operations (redis-py compatibility)."""
+        self._execute()
+        return []