mi-crow 0.1.1.post13__py3-none-any.whl → 0.1.1.post15__py3-none-any.whl

mi_crow/datasets/base_dataset.py

@@ -454,7 +454,7 @@ class BaseDataset(ABC):
         limit: Optional[int] = None,
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> "BaseDataset":
         """
         Load dataset from HuggingFace Hub.
@@ -534,7 +534,7 @@ class BaseDataset(ABC):
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
         drop_na_columns: Optional[List[str]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> "BaseDataset":
         """
         Load dataset from CSV file.
@@ -593,7 +593,7 @@ class BaseDataset(ABC):
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
         drop_na_columns: Optional[List[str]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> "BaseDataset":
         """
         Load dataset from JSON or JSONL file.

mi_crow/datasets/classification_dataset.py

@@ -368,7 +368,7 @@ class ClassificationDataset(BaseDataset):
         stratify_seed: Optional[int] = None,
         streaming: Optional[bool] = None,
         drop_na: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "ClassificationDataset":
         """
         Load classification dataset from HuggingFace Hub.
@@ -459,7 +459,7 @@ class ClassificationDataset(BaseDataset):
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
         drop_na: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "ClassificationDataset":
         """
         Load classification dataset from CSV file.
@@ -519,7 +519,7 @@ class ClassificationDataset(BaseDataset):
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
         drop_na: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "ClassificationDataset":
         """
         Load classification dataset from JSON/JSONL file.

mi_crow/datasets/text_dataset.py

@@ -233,7 +233,7 @@ class TextDataset(BaseDataset):
         stratify_seed: Optional[int] = None,
         streaming: Optional[bool] = None,
         drop_na: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "TextDataset":
         """
         Load text dataset from HuggingFace Hub.
@@ -312,7 +312,7 @@ class TextDataset(BaseDataset):
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
         drop_na: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "TextDataset":
         """
         Load text dataset from CSV file.
@@ -365,7 +365,7 @@ class TextDataset(BaseDataset):
         stratify_by: Optional[str] = None,
         stratify_seed: Optional[int] = None,
         drop_na: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "TextDataset":
         """
         Load text dataset from JSON/JSONL file.

mi_crow/language_model/language_model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections import defaultdict
 from pathlib import Path
-from typing import Sequence, Any, Dict, List, TYPE_CHECKING, Set
+from typing import Sequence, Any, Dict, List, TYPE_CHECKING, Set, Tuple
 
 import torch
 from torch import nn, Tensor
@@ -134,7 +134,7 @@ class LanguageModel:
         """Set the store instance."""
         self.context.store = value
 
-    def tokenize(self, texts: Sequence[str], **kwargs: Any):
+    def tokenize(self, texts: Sequence[str], **kwargs: Any) -> Any:
         """
         Tokenize texts using the language model tokenizer.
         
@@ -154,7 +154,7 @@ class LanguageModel:
             autocast: bool = True,
             autocast_dtype: torch.dtype | None = None,
             with_controllers: bool = True,
-    ):
+    ) -> Tuple[Any, Any]:
         """
         Run forward pass on texts.
         

mi_crow/language_model/layers.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Callable, TYPE_CHECKING
+from typing import Dict, List, Callable, TYPE_CHECKING, Any
 
 from torch import nn
 
@@ -121,7 +121,7 @@ class LanguageModelLayers:
             layer_signature: str | int,
             hook: Callable,
             hook_args: dict = None
-    ):
+    ) -> Any:
         """
         Register a forward hook directly on a layer.
         
@@ -141,7 +141,7 @@ class LanguageModelLayers:
             layer_signature: str | int,
             hook: Callable,
             hook_args: dict = None
-    ):
+    ) -> Any:
         """
         Register a pre-forward hook directly on a layer.
         

mi_crow/mechanistic/sae/modules/topk_sae.py

@@ -26,6 +26,8 @@ class TopKSaeTrainingConfig(SaeTrainingConfig):
     
     Args:
         k: Number of top activations to keep (required for TopK SAE training)
+    
+    Note:
         All other parameters are inherited from SaeTrainingConfig.
     
     Attributes:

mi_crow/mechanistic/sae/training/__init__.py

@@ -0,0 +1,6 @@
+"""Training utilities for SAE models."""
+
+from mi_crow.mechanistic.sae.training.wandb_logger import WandbLogger
+
+__all__ = ["WandbLogger"]
+

{mi_crow-0.1.1.post13.dist-info → mi_crow-0.1.1.post15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mi-crow
-Version: 0.1.1.post13
+Version: 0.1.1.post15
 Summary: Engineer Thesis: Explaining and modifying LLM responses using SAE and concepts.
 Author-email: Hubert Kowalski <your.email@example.com>, Adam Kaniasty <adam.kaniasty@gmail.com>
 Requires-Python: >=3.10

{mi_crow-0.1.1.post13.dist-info → mi_crow-0.1.1.post15.dist-info}/RECORD

@@ -1,10 +1,10 @@
 mi_crow/__init__.py,sha256=J7aXVlAicbjvk5630rhDxx0ATsvZnihud5u_aQpAwY8,487
 mi_crow/utils.py,sha256=LTfh2Ep87lAgPBaZkrQPP9caXFJoS9zUxu4qFuV4kzM,1549
 mi_crow/datasets/__init__.py,sha256=lCAc3nFlvoERrBPAan6C9YFmDx86W2gbIAy267Rb2Sk,349
-mi_crow/datasets/base_dataset.py,sha256=vYx-oj3jVhLZD1-xGSO4K4ZIsQtYpHP5zHmg7jd4FE0,22512
-mi_crow/datasets/classification_dataset.py,sha256=nL_xndJHyf8hlLxKBe_ZO2YLYsXQjGyeY6csqGTTzEY,21706
+mi_crow/datasets/base_dataset.py,sha256=fSs63o0LDNXt1-i0HmKTxaP33CMcS6rWJ3C4vpB8qR0,22527
+mi_crow/datasets/classification_dataset.py,sha256=c2peu0kWVY3N8XlEi9GjRuE6MmVV8f8vz8gtPmxI70s,21721
 mi_crow/datasets/loading_strategy.py,sha256=17VM3Td8lqDllGIx9DHI6WiXmSKKQHDHbfe4ZeM8ATA,1206
-mi_crow/datasets/text_dataset.py,sha256=5FzHWkMWWK0yP69O48S3fUj5KgHb8qo3mkvvZihHFuU,16781
+mi_crow/datasets/text_dataset.py,sha256=dk5RzWy-T7ssQNwv84FrgEZ83zuDgpq91XgUfbwvk8c,16796
 mi_crow/hooks/__init__.py,sha256=KYy5qcbEpnJceNH86ofy43Suu_36QXjj0HYl79rVyls,693
 mi_crow/hooks/controller.py,sha256=eo8LMERORXYUjH4-_R6DHk5JKN6O8SW6PlnuBFrlNqg,6063
 mi_crow/hooks/detector.py,sha256=Bj3xz56cSgRvbcoQBsHIdlJdf0dtgVLw3l1pOSRvRAg,3114
@@ -22,8 +22,8 @@ mi_crow/language_model/contracts.py,sha256=6ij7rzJcpSAKgYx-fiefg0Fi8TsFugaM5Ilih
 mi_crow/language_model/hook_metadata.py,sha256=GACZjZUneo2l5j7DCFycLAunTm0etdMQ2YB_xgueUuk,1394
 mi_crow/language_model/inference.py,sha256=-Kpm85jM8y6-GyDgrvIczitBIwGh8grJP8aYuXsLV-g,19082
 mi_crow/language_model/initialization.py,sha256=e_Vkk-p9KWRt6-Hmkm6I29dTf20jzEAyNF9CG4nc48M,3704
-mi_crow/language_model/language_model.py,sha256=a6CcklVA65oYtFxGXiwQrOKMPZj6eb7LOiT1zJ5-guo,13965
-mi_crow/language_model/layers.py,sha256=1yExHodMyqr_Yk4W-2HiSGnRs2sYOA7swsxI8u0Uvfk,15914
+mi_crow/language_model/language_model.py,sha256=b4KniweHXauPHPgHba9cAlLK86z2atfeoetnG-UHoSo,13998
+mi_crow/language_model/layers.py,sha256=7RdVU1mXt5oy2OhzNoKDeD9Omm899ZC-BJqpcTfxP2w,15933
 mi_crow/language_model/persistence.py,sha256=9wQE6tRvLg7BgdLlkKRTOfRwXb5Q0LsEgg8B9J7Yos0,5881
 mi_crow/language_model/tokenizer.py,sha256=uZbMDVNnzu8WZINUaR1tLFXiuk9V5pAoahwnJOUvEuE,7379
 mi_crow/language_model/utils.py,sha256=5Y7scRvvudUjKDV8QPhC3HAc2S-dCuqbm6xEjRr0fRM,2630
@@ -39,13 +39,14 @@ mi_crow/mechanistic/sae/concepts/concept_models.py,sha256=HGyPoMSmj8CAg9joIa6fV_
 mi_crow/mechanistic/sae/concepts/input_tracker.py,sha256=kIiqt7guv_-9-UPYtefAFJbHkWtAS_mnqYVvRU4eb2o,1890
 mi_crow/mechanistic/sae/modules/__init__.py,sha256=e0lkCALQZcJN7KpYyTtXx3OD2NhBxV_kOZLLJ6EWaTE,243
 mi_crow/mechanistic/sae/modules/l1_sae.py,sha256=qqw0iTWLSmWAlz5kgfw_mex8LeecFWM1FobyUteMqmM,15388
-mi_crow/mechanistic/sae/modules/topk_sae.py,sha256=pK_ajKTQb0wGAftzb6AE5ZZthV3aFLr6G3avOVclSHE,17313
+mi_crow/mechanistic/sae/modules/topk_sae.py,sha256=TIueJ4ftJm0XfU2UyHm0n3bIwL-W1RUgnpYym4xBxpg,17328
+mi_crow/mechanistic/sae/training/__init__.py,sha256=5flCJVkOyKizY0FZy1OP5v0EI6bPEayunpnUPp82a6s,140
 mi_crow/mechanistic/sae/training/wandb_logger.py,sha256=YlSJd5CaNa35RmIgf1FD_gSEDyhGRa2UdHo_Ofrplos,8558
 mi_crow/store/__init__.py,sha256=DrYTpdgzrRzjHm9bigy-GiP0BGxzjmD3-lJCthtgxbE,123
 mi_crow/store/local_store.py,sha256=XmguFvdrUi6NHzvV_bLaDJzpk5KWU_-ObkzhICcLu6g,17216
 mi_crow/store/store.py,sha256=VuDe9Git0glND3TTHh0zhDJNxdQY3dCp0cURhApYQbU,9334
 mi_crow/store/store_dataloader.py,sha256=UkZhHCOTg56ozomPtU9vHBhxIMOPcOiyfMqiAxgqtQs,4341
-mi_crow-0.1.1.post13.dist-info/METADATA,sha256=5KEjRwTvthwSs5Jed9Nr0TCKDpfzvLicbVKKm6KkBnQ,6584
-mi_crow-0.1.1.post13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mi_crow-0.1.1.post13.dist-info/top_level.txt,sha256=DTuNo2VWgrH6jQKY19NciReSpLwGKKIRzJ3WbpspLlE,8
-mi_crow-0.1.1.post13.dist-info/RECORD,,
+mi_crow-0.1.1.post15.dist-info/METADATA,sha256=lsLbGfYLFGEQuQALB9AOqOGyZovTk7Of7gvbKuI5fOY,6584
+mi_crow-0.1.1.post15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mi_crow-0.1.1.post15.dist-info/top_level.txt,sha256=DTuNo2VWgrH6jQKY19NciReSpLwGKKIRzJ3WbpspLlE,8
+mi_crow-0.1.1.post15.dist-info/RECORD,,