Merge branch 'fix_types' of https://github.com/EleutherAI/delphi into neighbour_latents

Author: SrGonao

delphi/latents/samplers.py

@@ -1,8 +1,6 @@
 import random
 from collections import deque
-from typing import Literal, cast
-
-from torchtyping import TensorType
+from typing import Literal
 
 from ..config import ExperimentConfig
 from ..logger import logger
@@ -121,10 +119,9 @@ def train(
             selected_examples = []
             for quantile in selected_examples_quantiles:
                 for example in quantile:
-                    example.normalized_activations = cast(
-                        TensorType["seq"],
-                        (example.activations * 10 / max_activation).floor(),
-                    )
+                    example.normalized_activations = (
+                        example.activations * 10 / max_activation
+                    ).floor()
                 selected_examples.extend(quantile)
             return selected_examples
 

delphi/utils.py

@@ -1,6 +1,7 @@
 from typing import Any, Type, TypeVar, cast
 
-from torchtyping import TensorType
+from jaxtyping import Float
+from torch import Tensor
 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
 
 
@@ -26,7 +27,7 @@ def load_tokenized_data(
     )
     tokens_ds = tokens_ds.shuffle(seed)
 
-    tokens = cast(TensorType["batch", "seq"], tokens_ds["input_ids"])
+    tokens = cast(Float[Tensor, "batch seq"], tokens_ds["input_ids"])
 
     return tokens