Merge pull request #114 from EleutherAI/frequency_estimation

Add per token and per context frequency to LatentRecord

Author: SrGonao

delphi/latents/constructors.py

@@ -159,6 +159,9 @@ def constructor(
     non_active_indices = mask.nonzero(as_tuple=False).squeeze()
     activations = activation_data.activations
 
+    # per context frequency
+    record.per_context_frequency = len(unique_batch_pos) / n_windows
+
     # Add activation examples to the record in place
     token_windows, act_windows = pool_max_activation_windows(
         activations=activations,

delphi/latents/latents.py

@@ -146,6 +146,13 @@ class LatentRecord:
     extra_examples: Optional[list[Example]] = None
     """Extra examples to include in the record."""
 
+    per_token_frequency: float = 0.0
+    """Frequency of the latent. Number of activations per total number of tokens."""
+
+    per_context_frequency: float = 0.0
+    """Frequency of the latent. Number of activations in a context per total
+    number of contexts."""
+
     @property
     def max_activation(self) -> float:
         """

delphi/latents/loader.py

@@ -378,6 +378,14 @@ async def _aprocess_latent(self, latent_data: LatentData) -> LatentRecord | None
         if self.tokens is None:
             raise ValueError("Tokens are not loaded")
         record = LatentRecord(latent_data.latent)
+
+        # number of activations in the latent
+        n_active = len(latent_data.activation_data.activations)
+        # number of tokens in the latent
+        n_tokens = self.tokens.shape[1] * self.tokens.shape[0]
+        # frequency of the latent
+        record.per_token_frequency = n_active / n_tokens
+
         if self.neighbours is not None:
             record.set_neighbours(
                 self.neighbours[latent_data.module][