feat: replaces print with logging (#136)

* feat: replaces print with logging

* Change some infos to warnings in constructor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Goncalo Paulo <30472805+SrGonao@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

__init__.py

@@ -0,0 +1,5 @@
+# This logger is needed for running tests from the repo directory.
+# The actual package logger is in delphi/delphi/__init__.py
+import logging
+
+logger = logging.getLogger(__name__)

delphi/__init__.py

@@ -1 +1,5 @@
 __version__ = "0.0.2"
+
+import logging
+
+logger = logging.getLogger(__name__)

delphi/__main__.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import os
 from functools import partial
 from pathlib import Path
@@ -17,6 +18,7 @@
     PreTrainedTokenizerFast,
 )
 
+from delphi import logger
 from delphi.clients import Offline, OpenRouter
 from delphi.config import RunConfig
 from delphi.explainers import ContrastiveExplainer, DefaultExplainer, NoOpExplainer
@@ -450,6 +452,16 @@ async def run(
 
 
 if __name__ == "__main__":
+    # Configure logging for CLI usage
+    logger.setLevel(logging.INFO)
+    file_handler = logging.FileHandler("delphi.log")
+    file_handler.setLevel(logging.INFO)
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
     parser = ArgumentParser()
     parser.add_arguments(RunConfig, dest="run_cfg")
     args = parser.parse_args()

delphi/clients/offline.py

@@ -12,7 +12,8 @@
     destroy_model_parallel,
 )
 
-from ..logger import logger
+from delphi import logger
+
 from .client import Client, Response
 
 

delphi/clients/openrouter.py

@@ -3,7 +3,8 @@
 
 import httpx
 
-from ..logger import logger
+from delphi import logger
+
 from .client import Client, Response
 from .types import ChatFormatRequest
 
@@ -30,7 +31,7 @@ def __init__(
         self.temperature = temperature
         timeout_config = httpx.Timeout(5.0)
         self.client = httpx.AsyncClient(timeout=timeout_config)
-        print("WARNING: We currently don't support logprobs for OpenRouter")
+        logger.warning("We currently don't support logprobs for OpenRouter")
 
     def postprocess(self, response):
         response_json = response.json()

delphi/explainers/contrastive_explainer.py

@@ -60,15 +60,15 @@ async def __call__(self, record: LatentRecord) -> ExplainerResult:
                 response_text = response
             explanation = self.parse_explanation(response_text)
             if self.verbose:
-                from ..logger import logger
+                from delphi import logger
 
                 logger.info(f"Explanation: {explanation}")
                 logger.info(f"Messages: {messages[-1]['content']}")
                 logger.info(f"Response: {response}")
 
             return ExplainerResult(record=record, explanation=explanation)
         except Exception as e:
-            from ..logger import logger
+            from delphi import logger
 
             logger.error(f"Explanation parsing failed: {repr(e)}")
             return ExplainerResult(

delphi/explainers/explainer.py

@@ -8,9 +8,10 @@
 
 import aiofiles
 
+from delphi import logger
+
 from ..clients.client import Client, Response
 from ..latents.latents import ActivatingExample, LatentRecord
-from ..logger import logger
 
 
 class ExplainerResult(NamedTuple):
@@ -127,7 +128,7 @@ async def explanation_loader(
             explanation = json.loads(await f.read())
         return ExplainerResult(record=record, explanation=explanation)
     except FileNotFoundError:
-        print(f"No explanation found for {record.latent}")
+        logger.info(f"No explanation found for {record.latent}")
         return ExplainerResult(record=record, explanation="No explanation found")
 
 

delphi/latents/cache.py

@@ -12,6 +12,7 @@
 from tqdm import tqdm
 from transformers import PreTrainedModel
 
+from delphi import logger
 from delphi.config import CacheConfig
 from delphi.latents.collect_activations import collect_activations
 
@@ -298,7 +299,7 @@ def run(self, n_tokens: int, tokens: token_tensor_type):
                 pbar.update(1)
                 pbar.set_postfix({"Total Tokens": f"{total_tokens:,}"})
 
-        print(f"Total tokens processed: {total_tokens:,}")
+        logger.info(f"Total tokens processed: {total_tokens:,}")
         self.cache.save()
         self.save_firing_counts()
 
@@ -374,8 +375,8 @@ def save_splits(self, n_splits: int, save_dir: Path, save_tokens: bool = True):
                     masked_locations = masked_locations.astype(np.uint16)
                 else:
                     masked_locations = masked_locations.astype(np.uint32)
-                    print(
-                        "Warning: Increasing the number of splits might reduce the"
+                    logger.warning(
+                        "Increasing the number of splits might reduce the"
                         "memory usage of the cache."
                     )
 
@@ -399,10 +400,10 @@ def generate_statistics_cache(self):
         to the console.
         """
         assert self.width is not None, "Width must be set before generating statistics"
-        print("Feature statistics:")
+        logger.info("Feature statistics:")
         # Token frequency
         for module_path in self.cache.latent_locations.keys():
-            print(f"# Module: {module_path}")
+            logger.info(f"# Module: {module_path}")
             generate_statistics_cache(
                 self.cache.tokens[module_path],
                 self.cache.latent_locations[module_path],
@@ -493,7 +494,7 @@ def generate_statistics_cache(
     num_alive = counts.shape[0]
     fraction_alive = num_alive / width
     if verbose:
-        print(f"Fraction of latents alive: {fraction_alive:%}")
+        logger.info(f"Fraction of latents alive: {fraction_alive:%}")
     # Compute densities of latents
     densities = counts / total_n_tokens
 
@@ -502,8 +503,12 @@ def generate_statistics_cache(
     # How many fired more than 10% of the time
     ten_percent = (densities > 0.1).sum() / width
     if verbose:
-        print(f"Fraction of latents fired more than 1% of the time: {one_percent:%}")
-        print(f"Fraction of latents fired more than 10% of the time: {ten_percent:%}")
+        logger.info(
+            f"Fraction of latents fired more than 1% of the time: {one_percent:%}"
+        )
+        logger.info(
+            f"Fraction of latents fired more than 10% of the time: {ten_percent:%}"
+        )
     # Try to estimate simple feature frequency
     split_indices = torch.cumsum(counts, dim=0)
     activation_splits = torch.tensor_split(sorted_activations, split_indices[:-1])
@@ -525,8 +530,10 @@ def generate_statistics_cache(
     single_token_fraction = maybe_single_token_features / num_alive
     strong_token_fraction = num_single_token_features / num_alive
     if verbose:
-        print(f"Fraction of weak single token latents: {single_token_fraction:%}")
-        print(f"Fraction of strong single token latents: {strong_token_fraction:%}")
+        logger.info(f"Fraction of weak single token latents: {single_token_fraction:%}")
+        logger.info(
+            f"Fraction of strong single token latents: {strong_token_fraction:%}"
+        )
 
     return CacheStatistics(
         frac_alive=float(fraction_alive),

delphi/latents/constructors.py

@@ -11,6 +11,8 @@
 from torch import Tensor
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
+from delphi import logger
+
 from ..config import ConstructorConfig
 from .latents import (
     ActivatingExample,
@@ -25,7 +27,7 @@
 def get_model(name: str, device: str = "cuda") -> SentenceTransformer:
     global model_cache
     if (name, device) not in model_cache:
-        print(f"Loading model {name} on device {device}")
+        logger.info(f"Loading model {name} on device {device}")
         model_cache[(name, device)] = SentenceTransformer(name, device=device)
     return model_cache[(name, device)]
 
@@ -284,7 +286,9 @@ def constructor(
         for toks, acts in zip(token_windows, act_windows)
     ]
     if len(record.examples) < min_examples:
-        print(f"Not enough examples to explain the latent: {len(record.examples)}")
+        logger.warning(
+            f"Not enough examples to explain the latent: {len(record.examples)}"
+        )
         # Not enough examples to explain the latent
         return None
 
@@ -404,7 +408,7 @@ def faiss_non_activation_windows(
 
     # Check if we have enough non-activating examples
     if available_indices.numel() < n_not_active:
-        print("Not enough non-activating examples available")
+        logger.warning("Not enough non-activating examples available")
         return []
 
     # Reshape tokens to get context windows
@@ -426,7 +430,7 @@ def faiss_non_activation_windows(
     ]
 
     if not activating_texts:
-        print("No activating examples available")
+        logger.warning("No activating examples available")
         return []
 
     # Create unique cache keys for both activating and non-activating texts
@@ -451,17 +455,17 @@ def faiss_non_activation_windows(
     if cache_enabled and non_activating_cache_file.exists():
         try:
             index = faiss.read_index(str(non_activating_cache_file), faiss.IO_FLAG_MMAP)
-            print(f"Loaded non-activating index from {non_activating_cache_file}")
+            logger.info(f"Loaded non-activating index from {non_activating_cache_file}")
         except Exception as e:
-            print(f"Error loading cached embeddings: {repr(e)}")
+            logger.warning(f"Error loading cached embeddings: {repr(e)}")
 
     if index is None:
-        print("Decoding non-activating tokens...")
+        logger.info("Decoding non-activating tokens...")
         non_activating_texts = [
             "".join(tokenizer.batch_decode(tokens)) for tokens in non_activating_tokens
         ]
 
-        print("Computing non-activating embeddings...")
+        logger.info("Computing non-activating embeddings...")
         non_activating_embeddings = get_model(embedding_model).encode(
             non_activating_texts, show_progress_bar=False
         )
@@ -472,26 +476,30 @@ def faiss_non_activation_windows(
         if cache_enabled:
             os.makedirs(cache_path, exist_ok=True)
             faiss.write_index(index, str(non_activating_cache_file))
-            print(f"Cached non-activating embeddings to {non_activating_cache_file}")
+            logger.info(
+                f"Cached non-activating embeddings to {non_activating_cache_file}"
+            )
 
     activating_embeddings = None
     if cache_enabled and activating_cache_file.exists():
         try:
             activating_embeddings = np.load(activating_cache_file)
-            print(f"Loaded cached activating embeddings from {activating_cache_file}")
+            logger.info(
+                f"Loaded cached activating embeddings from {activating_cache_file}"
+            )
         except Exception as e:
-            print(f"Error loading cached embeddings: {repr(e)}")
+            logger.warning(f"Error loading cached embeddings: {repr(e)}")
     # Compute embeddings for activating examples if not cached
     if activating_embeddings is None:
-        print("Computing activating embeddings...")
+        logger.info("Computing activating embeddings...")
         activating_embeddings = get_model(embedding_model).encode(
             activating_texts, show_progress_bar=False
         )
         # Cache the embeddings
         if cache_enabled:
             os.makedirs(cache_path, exist_ok=True)
             np.save(activating_cache_file, activating_embeddings)
-            print(f"Cached activating embeddings to {activating_cache_file}")
+            logger.info(f"Cached activating embeddings to {activating_cache_file}")
 
     # Search for the nearest neighbors to each activating example
     collected_indices = set()
@@ -618,7 +626,9 @@ def neighbour_non_activation_windows(
         )
         number_examples += examples_used
     if len(all_examples) == 0:
-        print("No examples found, falling back to random non-activating examples")
+        logger.warning(
+            "No examples found, falling back to random non-activating examples"
+        )
         non_active_indices = not_active_mask.nonzero(as_tuple=False).squeeze()
 
         return random_non_activating_windows(
@@ -655,7 +665,7 @@ def random_non_activating_windows(
     # If this happens it means that the latent is active in every window,
     # so it is a bad latent
     if available_indices.numel() < n_not_active:
-        print("No available randomly sampled non-activating sequences")
+        logger.warning("No available randomly sampled non-activating sequences")
         return []
     else:
         random_indices = torch.randint(

delphi/latents/samplers.py

@@ -6,8 +6,9 @@
     PreTrainedTokenizerFast,
 )
 
+from delphi import logger
+
 from ..config import SamplerConfig
-from ..logger import logger
 from .latents import ActivatingExample, LatentRecord