Merge pull request #54 from JadenFiotto-Kaufman/dev

Dev

src/nnsight/__init__.py

@@ -14,6 +14,10 @@
 from .models.LanguageModel import LanguageModel
 from .module import Module
 from .patching import Patch, Patcher
+from .logger import logger
+
+logger.disabled = not CONFIG.APP.LOGGING
+
 
 # Below do default patching:
 DEFAULT_PATCHER = Patcher()

src/nnsight/config.yaml

@@ -1,2 +1,4 @@
 API:
-  HOST: ndif.dev
+  HOST: ndif.dev
+APP:
+  LOGGING: False

src/nnsight/models/LanguageModel.py

@@ -76,6 +76,7 @@ def _tokenize(
             torch.Tensor,
             Dict[str, Any],
         ],
+        **kwargs,
     ):
         if isinstance(inputs, BatchEncoding):
             return inputs
@@ -93,7 +94,7 @@ def _tokenize(
 
             return self.tokenizer.pad(inputs, return_tensors="pt")
 
-        return self.tokenizer(inputs, return_tensors="pt", padding=True)
+        return self.tokenizer(inputs, return_tensors="pt", padding=True, **kwargs)
 
     def _prepare_inputs(
         self,
@@ -111,23 +112,23 @@ def _prepare_inputs(
         **kwargs,
     ) -> BatchEncoding:
         if isinstance(inputs, dict):
-            _inputs = self._tokenize(inputs["input_ids"])
-
-            for ai, attn_mask in enumerate(inputs['attention_mask']):
+            _inputs = self._tokenize(inputs["input_ids"], **kwargs)
 
-                _inputs['attention_mask'][ai, -len(attn_mask):] = attn_mask
+            if "attention_mask" in inputs:
+                for ai, attn_mask in enumerate(inputs["attention_mask"]):
+                    _inputs["attention_mask"][ai, -len(attn_mask) :] = attn_mask
 
             if "labels" in inputs:
-                labels = self._tokenize(inputs["labels"])
-                labels = self._tokenize(labels)
+                labels = self._tokenize(inputs["labels"], **kwargs)
+
                 _inputs["labels"] = labels["input_ids"]
 
             return _inputs
 
-        inputs = self._tokenize(inputs)
+        inputs = self._tokenize(inputs, **kwargs)
 
         if labels is not None:
-            labels = self._tokenize(labels)
+            labels = self._tokenize(labels, **kwargs)
 
             inputs["labels"] = labels["input_ids"]
 

src/nnsight/pydantics/Config.py

@@ -5,5 +5,10 @@ class ApiConfigModel(BaseModel):
     HOST: str
 
 
+class AppConfigModel(BaseModel):
+    LOGGING: bool
+
+
 class ConfigModel(BaseModel):
     API: ApiConfigModel
+    APP: AppConfigModel