Extend documentation: Comprehensive docstring enhancement across all modules

README.md

@@ -1,54 +1,220 @@
 # EasyRoutine
 
-This is just a simple collection of routines that I use frequently. I have found that I often need to do the same things over and over again, so I have created this repository to store them. I hope you find them useful.
+A comprehensive Python toolkit for AI model interpretability, analysis, and efficient inference. EasyRoutine provides researchers and practitioners with powerful tools for understanding transformer models, extracting internal activations, and performing mechanistic interpretability studies.
 
-## Installation
+## 🚀 Key Features
 
+- **🔍 Mechanistic Interpretability**: Deep analysis of transformer model internals
+  - Extract activations from any model component (residual streams, attention layers, MLPs)
+  - Support for attention pattern analysis and head-specific extractions
+  - Intervention capabilities for ablation studies and activation patching
+  - Works with both language models and vision-language models
 
-## Interpretability
-The interpretability module contains wrapper of huggingface LLM/VLM that help to perform interpretability tasks on the model. Currently, it supports:
-- Extract activations of any component of the model
-- Perform ablation study on the model during inference
-- Perform activation patching on the model during inference
+- **⚡ High-Performance Inference**: Optimized model inference with multiple backends
+  - VLLM integration for high-throughput inference
+  - Multi-GPU support and memory optimization
+  - Configurable generation parameters and chat templates
+
+- **📊 Smart Progress Tracking**: Adaptive progress bars for any environment
+  - Rich progress bars for interactive use
+  - Clean text logging for batch jobs (SLURM, PBS, etc.)
+  - Automatic environment detection
+
+- **🛠️ Essential Utilities**: Common tools for ML workflows
+  - Robust logging with multiple output formats
+  - File system navigation helpers
+  - Memory management utilities
+
+## 📦 Installation
+
+```bash
+pip install easyroutine
+```
+
+For development installation:
+```bash
+git clone https://github.com/francescortu/easyroutine.git
+cd easyroutine
+pip install -e .
+```
+
+## 🔍 Interpretability - Quick Start
+
+### Basic Activation Extraction
 
-### Simple Tutorial
 ```python
-# First we need to import the HookedModel and the config classes
-from easyroutine.interpretability import HookedModel,  ExtractionConfig
+from easyroutine.interpretability import HookedModel, ExtractionConfig
+from easyroutine.console import progress
 
-hooked_model = HookedModel.from_pretrained(
-    model_name="mistral-community/pixtral-12b", # the model name
-    device_map = "auto"
+# Load any Hugging Face transformer model
+model = HookedModel.from_pretrained(
+    model_name="gpt2",  # or any HF model
+    device_map="auto"
 )
 
-# Now let's define a simple dataset
-dataset = [
-    "This is a test",
-    "This is another test"
-]
+# Prepare your data
+texts = ["Hello, world!", "How are you today?"]
+tokenizer = model.get_tokenizer()
+dataset = [tokenizer(text, return_tensors="pt") for text in texts]
+
+# Configure what activations to extract
+config = ExtractionConfig(
+    extract_resid_out=True,        # Residual stream outputs
+    extract_attn_pattern=True,     # Attention patterns  
+    extract_mlp_out=True,          # MLP layer outputs
+    save_input_ids=True            # Keep track of tokens
+)
 
-tokenizer = hooked_model.get_tokenizer()
+# Extract activations with progress tracking
+cache = model.extract_cache(
+    progress(dataset, description="Extracting activations"),
+    target_token_positions=["last"],  # Focus on final token
+    extraction_config=config
+)
 
-dataset = tokenizer(dataset, padding=True, truncation=True, return_tensors="pt") 
+# Access extracted data
+residual_activations = cache["resid_out_0"]  # Layer 0 residual outputs
+attention_patterns = cache["attn_pattern_0_5"]  # Layer 0, Head 5 attention
+print(f"Extracted activations: {list(cache.keys())}")
+```
 
-cache = hooked_model.extract_cache(
-    dataset,
-    target_token_positions = ["last"],
-    extraction_config = ExtractionConfig(
-        extract_resid_out = True
+### Advanced: Intervention Studies
+
+```python
+from easyroutine.interpretability import Intervention
+
+# Define interventions for causal analysis
+interventions = [
+    Intervention(
+        component="resid_out_5",  # Target layer 5 residual stream
+        intervention_type="ablation",  # Zero out activations
+        positions=["last"]  # Only affect last token
     )
+]
+
+# Run model with interventions
+cache_with_intervention = model.extract_cache(
+    dataset,
+    target_token_positions=["last"],
+    extraction_config=config,
+    interventions=interventions
+)
+
+# Compare original vs. intervened activations
+original_logits = cache["logits"]
+intervened_logits = cache_with_intervention["logits"] 
+effect = original_logits - intervened_logits
+```
+
+### Vision-Language Models
+
+```python
+# Works seamlessly with VLMs like LLaVA, Pixtral, etc.
+vlm = HookedModel.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+
+# Process multimodal inputs
+processor = vlm.get_processor()
+inputs = processor(images=image, text="What do you see?", return_tensors="pt")
+
+# Extract activations from multimodal processing
+cache = vlm.extract_cache(
+    [inputs],
+    target_token_positions=["last-image", "end-image"],
+    extraction_config=ExtractionConfig(extract_resid_out=True)
+)
+```
+
+## ⚡ High-Performance Inference
+
+```python
+from easyroutine.inference import VLLMInferenceModel
+
+# Initialize high-performance inference engine
+model = VLLMInferenceModel.init_model(
+    model_name="microsoft/DialoGPT-large",
+    n_gpus=2,
+    dtype="bfloat16"
+)
+
+# Generate responses
+response = model.generate("Hello, how can I help you today?")
+
+# Multi-turn conversations with chat templates
+chat_history = []
+chat_history = model.append_with_chat_template(
+    "What is machine learning?",
+    role="user",
+    chat_history=chat_history
+)
+response = model.generate_with_chat_template(chat_history)
+```
+
+## 🛠️ Utilities & Logging
+
+```python
+from easyroutine.logger import setup_logging, logger
+from easyroutine import path_to_parents
+from easyroutine.console import progress
+
+# Flexible logging setup
+setup_logging(
+    level="INFO",
+    file="experiment.log", 
+    console=True,
+    fmt="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 
-````
+logger.info("Starting experiment...")
+
+# Convenient navigation helpers
+path_to_parents(2)  # Go up 2 directory levels
+
+# Smart progress tracking that adapts to your environment
+for item in progress(large_dataset, description="Processing data"):
+    result = expensive_computation(item)
+
+# In interactive environments: rich progress bar with ETA
+# In batch jobs (SLURM, etc.): clean timestamped logging
+# [2024-01-15 10:30:15] Processing data: 1500/10000 (15.0%) - Elapsed: 2.3m, Remaining: 13.1m
+```
+
+## 📖 Documentation
+
+- **API Reference**: Comprehensive docstrings for all functions and classes
+- **Examples**: Jupyter notebooks with common use cases
+- **Tutorials**: Step-by-step guides for interpretability workflows
+
+## 🤝 Contributing
+
+We welcome contributions! Please see our contributing guidelines for details on:
+- Code style and testing requirements
+- Documentation standards  
+- How to submit issues and pull requests
+
+## 📄 License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## 🔬 Citation
+
+If you use EasyRoutine in your research, please cite:
 
+```bibtex
+@software{easyroutine2024,
+  title={EasyRoutine: A Toolkit for AI Model Interpretability},
+  author={Francesco Ortu},
+  year={2024},
+  url={https://github.com/francescortu/easyroutine}
+}
+```
 
+---
 
 ### Development
-For publish the package push a commit with the flag:
-  - `[patch]`: x.x.7 -> x.x.8
-  - `[minor]`: x.7.x -> x.8.0
-  - `[major]`: 2.x.x -> 3.0.0
 
-Example commit: `fix multiple bus [patch]`
+For publishing new versions, use semantic version tags in commit messages:
+- `[patch]`: Bug fixes (x.x.7 → x.x.8) 
+- `[minor]`: New features (x.7.x → x.8.0)
+- `[major]`: Breaking changes (2.x.x → 3.0.0)
 
--
+Example: `git commit -m "Add support for Gemma models [minor]"`

easyroutine/__init__.py

@@ -1,2 +1,79 @@
+"""
+EasyRoutine: A comprehensive toolkit for AI model interpretability and analysis.
+
+EasyRoutine is a powerful Python package that provides a collection of utilities
+and tools for working with machine learning models, with a particular focus on
+mechanistic interpretability of transformer models. It offers a unified interface
+for model analysis, activation extraction, intervention studies, and more.
+
+Key Features:
+    - 🔍 Mechanistic Interpretability: Comprehensive tools for analyzing transformer
+      model internals, including activation extraction, attention pattern analysis,
+      and intervention studies
+    - 🚀 High-Performance Inference: Optimized inference backends with support for
+      VLLM and custom implementations
+    - 📊 Progress Tracking: Adaptive progress bars that work in both interactive
+      and batch environments
+    - 🛠️ Utility Functions: Common utilities for file system navigation and
+      robust logging
+    - 🔧 Extensible Architecture: Modular design allowing easy extension and
+      customization
+
+Main Modules:
+    interpretability: Core functionality for transformer model analysis
+        - HookedModel: Wrapper for extracting activations from any transformer
+        - ExtractionConfig: Fine-grained control over activation extraction
+        - ActivationCache: Efficient storage and manipulation of activations
+        - Intervention tools: Ablation studies and activation patching
+
+    inference: Model inference interfaces for various backends
+        - BaseInferenceModel: Abstract interface for inference implementations  
+        - VLLMInferenceModel: High-performance inference with VLLM backend
+        - Configuration classes for flexible inference setup
+
+    console: Progress tracking and console utilities
+        - progress(): tqdm-style progress bars with environment detection
+        - LoggingProgress: Text-based progress for batch environments
+
+    logger: Robust logging functionality
+        - Structured logging with multiple output options
+        - Level control and formatting customization
+        - Integration with rich for enhanced console output
+
+    utils: Common utility functions
+        - File system navigation helpers
+        - Path manipulation utilities
+
+Quick Start:
+    >>> import easyroutine
+    >>> from easyroutine.interpretability import HookedModel, ExtractionConfig
+    >>> from easyroutine.console import progress
+    >>> 
+    >>> # Load a model for interpretability analysis
+    >>> model = HookedModel.from_pretrained("gpt2")
+    >>> 
+    >>> # Configure what to extract
+    >>> config = ExtractionConfig(
+    ...     extract_resid_out=True,
+    ...     extract_attn_pattern=True
+    ... )
+    >>> 
+    >>> # Extract activations with progress tracking
+    >>> data = [{"input_ids": tokenized_inputs}]
+    >>> cache = model.extract_cache(
+    ...     progress(data, description="Extracting activations"),
+    ...     target_token_positions=["last"],
+    ...     extraction_config=config
+    ... )
+
+Installation:
+    The package can be installed via pip:
+    ```bash
+    pip install easyroutine
+    ```
+
+For more detailed documentation and examples, visit the project repository.
+"""
+
 from .utils import path_to_parents
 from .logger import logger