wisent 0.5.3__tar.gz → 0.5.5__tar.gz

{wisent-0.5.3 → wisent-0.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.3
+Version: 0.5.5
 Summary: Monitor and guard against harmful content in language models
 Home-page: https://github.com/yourusername/wisent-activation-guardrails
 Author: Wisent Team

wisent-0.5.5/wisent/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.5.5"

{wisent-0.5.3 → wisent-0.5.5}/wisent/core/__init__.py

@@ -1,9 +1,10 @@
 from .utils.device import empty_device_cache, preferred_dtype, resolve_default_device, resolve_device, resolve_torch_device
-from .steering import SteeringMethod, SteeringType
+# Note: SteeringMethod and SteeringType import temporarily disabled due to missing dependencies
+# from .steering import SteeringMethod, SteeringType
 
 __all__ = [
-    "SteeringMethod",
-    "SteeringType",
+    # "SteeringMethod",
+    # "SteeringType",
     "empty_device_cache",
     "preferred_dtype",
     "resolve_default_device",

{wisent-0.5.3 → wisent-0.5.5}/wisent/core/steering.py

@@ -7,8 +7,9 @@ from typing import Any, Dict, List, Optional, Union
 import torch
 import torch.nn.functional as F
 
-from wisent.core.activations import Activations
-from wisent.core.classifier.classifier import Classifier
+# Note: Activations and Classifier imports removed - these classes don't exist in the current package structure
+# from wisent.core.activations import Activations
+# from wisent.core.classifier.classifier import Classifier
 
 from .contrastive_pairs import ContrastivePairSet
 from .steering_method import CAA

wisent-0.5.5/wisent/core/utils/__init__.py

@@ -0,0 +1,21 @@
+from .device import (
+    DeviceKind,
+    empty_device_cache,
+    ensure_tensor_on_device,
+    move_module_to_preferred_device,
+    preferred_dtype,
+    resolve_default_device,
+    resolve_device,
+    resolve_torch_device,
+)
+
+__all__ = [
+    "DeviceKind",
+    "empty_device_cache",
+    "ensure_tensor_on_device",
+    "move_module_to_preferred_device",
+    "preferred_dtype",
+    "resolve_default_device",
+    "resolve_device",
+    "resolve_torch_device",
+]

wisent-0.5.5/wisent/core/utils/device.py

@@ -0,0 +1,56 @@
+"""Centralized torch device selection helpers."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Literal
+
+import torch
+
+DeviceKind = Literal["cuda", "mps", "cpu"]
+
+
+def _mps_available() -> bool:
+    return getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available()
+
+
+@lru_cache(maxsize=1)
+def resolve_default_device() -> DeviceKind:
+    if torch.cuda.is_available():
+        return "cuda"
+    if _mps_available():
+        return "mps"
+    return "cpu"
+
+
+def resolve_torch_device() -> torch.device:
+    return torch.device(resolve_default_device())
+
+
+def resolve_device(kind: DeviceKind | None = None) -> torch.device:
+    return torch.device(kind or resolve_default_device())
+
+
+def preferred_dtype(kind: DeviceKind | None = None) -> torch.dtype:
+    chosen = kind or resolve_default_device()
+    return torch.float16 if chosen in {"cuda", "mps"} else torch.float32
+
+
+def empty_device_cache(kind: DeviceKind | None = None) -> None:
+    chosen = kind or resolve_default_device()
+    if chosen == "cuda" and torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    elif chosen == "mps" and _mps_available():
+        try:
+            torch.mps.empty_cache()  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+
+
+def move_module_to_preferred_device(module: torch.nn.Module) -> torch.nn.Module:
+    return module.to(resolve_torch_device())
+
+
+def ensure_tensor_on_device(tensor: torch.Tensor) -> torch.Tensor:
+    target = resolve_torch_device()
+    return tensor.to(target) if tensor.device != target else tensor

{wisent-0.5.3 → wisent-0.5.5}/wisent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.3
+Version: 0.5.5
 Summary: Monitor and guard against harmful content in language models
 Home-page: https://github.com/yourusername/wisent-activation-guardrails
 Author: Wisent Team

{wisent-0.5.3 → wisent-0.5.5}/wisent.egg-info/SOURCES.txt

@@ -189,6 +189,8 @@ wisent/core/trainers/__init__.py
 wisent/core/trainers/steering_trainer.py
 wisent/core/trainers/core/__init__.py
 wisent/core/trainers/core/atoms.py
+wisent/core/utils/__init__.py
+wisent/core/utils/device.py
 wisent/opti/__init__.py
 wisent/opti/core/__init__.py
 wisent/opti/core/atoms.py

wisent-0.5.3/wisent/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.5.3"