Dev

pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
     "protobuf",
     "python-socketio[client]",
     "tokenizers>=0.13.0",
-    "pydantic>=2.4.0",
+    "pydantic>=2.9.0",
     "torch>=2.4.0",
     "sentencepiece",
     "torchvision",

src/nnsight/__init__.py

@@ -13,6 +13,13 @@
 from functools import wraps
 from typing import Dict, Union
 
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("nnsight")
+except PackageNotFoundError:
+    __version__ = "unknown version"
+
 import torch
 import yaml
 

src/nnsight/contexts/GraphBasedContext.py

@@ -307,6 +307,7 @@ class GlobalTracingContext(GraphBasedContext):
             global_patch(torch, "randn"),
             global_patch(torch, "randperm"),
             global_patch(torch, "zeros"),
+            global_patch(torch, "cat")
         ]
         + [
             global_patch_class(value)

src/nnsight/contexts/backends/RemoteBackend.py

@@ -215,7 +215,8 @@ def submit_request(self, request: "RequestModel") -> "ResponseModel":
 
         else:
 
-            raise Exception(response.reason)
+            msg = response.json()['detail']
+            raise ConnectionError(msg)
 
     def get_response(self) -> "ResponseModel":
         """Retrieves and handles the response object from the remote endpoint.

src/nnsight/envoy.py

@@ -142,12 +142,28 @@ def _set_tracer(self, tracer: Tracer, propagate=True):
         if propagate:
             for envoy in self._sub_envoys:
                 envoy._set_tracer(tracer, propagate=True)
+
+
+    def _tracing(self) -> bool:
+        """Whether or not tracing.
+
+        Returns:
+            bool: Is tracing.
+        """
+
+        try:
+
+            return self._tracer.graph.alive
+
+        except:
+
+            return False
 
     def _scanning(self) -> bool:
         """Whether or not in scanning mode. Checks the current Tracer's Invoker.
 
         Returns:
-            bool: _description_
+            bool: Is scanning.
         """
 
         try:
@@ -420,6 +436,9 @@ def __call__(
         Returns:
             InterventionProxy: Module call proxy.
         """
+
+        if not self._tracing():
+            return self._module(*args, **kwargs)
 
         if isinstance(self._tracer.backend, EditBackend):
             hook = True

src/nnsight/models/NNsightModel.py

@@ -119,6 +119,10 @@ def __init__(
             self.dispatch_model()
 
         logger.info(f"Initialized `{self._model_key}`")
+
+    def __call__(self, *args, **kwargs):
+
+        return self._envoy(*args, **kwargs)
 
     def trace(
         self,

src/nnsight/schema/format/functions.py

@@ -1,5 +1,6 @@
 import operator
 from inspect import getmembers, isbuiltin, isfunction, ismethod, ismethoddescriptor, isclass
+from typing import Callable
 
 import einops
 import torch
@@ -24,6 +25,15 @@ def get_function_name(fn, module_name=None):
 
     return f"{module_name}.{fn.__qualname__}"
 
+def update_function(function: str | Callable, new_function: Callable):
+
+    if not isinstance(function, str):
+
+        function = get_function_name(function)
+
+    new_function.__name__ = function
+
+    FUNCTIONS_WHITELIST[function] = new_function
 
 FUNCTIONS_WHITELIST = {}