Metadata-Version: 2.4
Name: nnez
Version: 0.1.0
Summary: Neural Network Easy Extraction - Simple LLM activation extraction in one line
Home-page: https://github.com/yourusername/nnez
Author: Your Name
Author-email: Paul Bogdan <paulcbogdan@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/yourusername/nnez
Project-URL: Documentation, https://nnez.readthedocs.io/
Project-URL: Repository, https://github.com/yourusername/nnez.git
Project-URL: Bug Tracker, https://github.com/yourusername/nnez/issues
Project-URL: Changelog, https://github.com/yourusername/nnez/blob/main/CHANGELOG.md
Keywords: llm,transformers,activation-extraction,interpretability,neural-networks,nlp,machine-learning,deep-learning,gpt,bert,probing,residual-stream
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=1.9.0
Requires-Dist: transformers>=4.20.0
Requires-Dist: nnsight>=0.2.0
Requires-Dist: numpy>=1.19.0
Requires-Dist: inflect>=6.0.0
Requires-Dist: typing-extensions>=4.0.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
Requires-Dist: black>=22.0.0; extra == "dev"
Requires-Dist: isort>=5.10.0; extra == "dev"
Requires-Dist: flake8>=4.0.0; extra == "dev"
Requires-Dist: mypy>=0.950; extra == "dev"
Requires-Dist: pre-commit>=2.17.0; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=4.5.0; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=1.0.0; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints>=1.18.0; extra == "docs"
Provides-Extra: examples
Requires-Dist: scipy>=1.7.0; extra == "examples"
Requires-Dist: matplotlib>=3.4.0; extra == "examples"
Requires-Dist: seaborn>=0.11.0; extra == "examples"
Requires-Dist: jupyter>=1.0.0; extra == "examples"
Requires-Dist: ipywidgets>=7.6.0; extra == "examples"
Dynamic: author
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-python

# nnez

**Neural Network Easy Extraction** - A lightweight Python package for extracting activation patterns from transformer language models with just a few lines of code.

[![PyPI version](https://badge.fury.io/py/nnez.svg)](https://badge.fury.io/py/nnez)
[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

## 🚀 Features

- 🎯 **Simple API** - Extract LLM activations with just one function call
- 🧠 **Multi-Model Support** - Works with GPT-2, GPT-Neo, LLaMA, BERT, and more
- 📊 **NumPy Output** - Get activations as NumPy arrays for easy analysis
- 🔄 **Flexible Layer Selection** - Extract from any combination of layers
- ⚡ **Automatic Caching** - Reuse loaded models for faster repeated extractions
- 🎨 **Bonus Grammar Tools** - Article detection and pluralization utilities using inflect

## 📦 Installation

```bash
pip install nnez
```

## 🎮 Quick Start

### Extract Activations from any Transformer LLM

```python
from nnez import get_activity_from_text

# Extract activations from specific layers of GPT-2
text = "The capital of France is Paris."
layers = [5, 10]  # Extract from layers 5 and 10

activations = get_activity_from_text(
    text=text,
    layers_list=layers,
    model_name="gpt2"
)

print(activations.shape)  # (2, 768) - 2 layers, 768 dimensions each
```

### Single Layer Extraction

```python
# Extract from a single layer (returns 1D array)
act = get_activity_from_text("Hello world!", 11)  # Layer 11 only
print(act.shape)  # (768,) - Single layer, flattened
```

### Batch Processing

```python
import numpy as np

texts = ["First text", "Second text", "Third text"]
all_activations = []

for text in texts:
    act = get_activity_from_text(text, [0, 6, 11])
    all_activations.append(act)

# Stack into 3D array: (num_texts, num_layers, hidden_size)
batch_activations = np.stack(all_activations)
print(batch_activations.shape)  # (3, 3, 768)
```

## 🔬 Advanced Usage

### Using Different Models

```python
# GPT-2 variants
act = get_activity_from_text(text, [10], model_name="gpt2-medium")

# GPT-Neo
act = get_activity_from_text(text, [0, 12], model_name="EleutherAI/gpt-neo-125M")

# BERT
act = get_activity_from_text(text, [6], model_name="bert-base-uncased")

# Any HuggingFace model
act = get_activity_from_text(text, [10], model_name="your-model-here")
```

### Layer Selection Strategies

```python
# First and last layers - for input/output representations
activations = get_activity_from_text(text, [0, 11])

# Every other layer - for efficient probing
activations = get_activity_from_text(text, list(range(0, 12, 2)))

# Middle layers only - for syntactic features
activations = get_activity_from_text(text, [4, 5, 6, 7])
```

### Similarity Analysis

```python
from scipy.spatial.distance import cosine

text1 = "The cat sat on the mat."
text2 = "The dog sat on the rug."

# Extract from same layers
act1 = get_activity_from_text(text1, [6, 10])
act2 = get_activity_from_text(text2, [6, 10])

# Compare layer-wise similarities
for i, layer in enumerate([6, 10]):
    similarity = 1 - cosine(act1[i], act2[i])
    print(f"Layer {layer} similarity: {similarity:.3f}")
```

## 🎁 Bonus: Grammar Utilities

The package includes grammar tools powered by the `inflect` library:

```python
from nnez.grammar import get_article, pluralize, quantify

# Smart article detection
get_article("hour")       # "an" (silent h)
get_article("university") # "a"  (y-sound)
get_article("FBI")        # "an" (eff-bee-eye)

# Pluralization
pluralize("child")        # "children"
pluralize("analysis")    # "analyses"
pluralize("octopus")     # "octopuses"

# Quantification
quantify(0, "cat")        # "no cats"
quantify(1, "child")      # "1 child"
quantify(3, "child")      # "3 children"
```

## 🛠️ API Reference

### Core Function

```python
get_activity_from_text(
    text: str,
    layers_list: Union[List[int], int],
    model_name: str = "gpt2",
    device: Optional[str] = None,
    verbose: bool = False,
    cache_model: bool = True
) -> np.ndarray
```

**Parameters:**
- `text`: Input text to analyze
- `layers_list`: Layer indices to extract from (list or single int)
- `model_name`: HuggingFace model identifier
- `device`: Device to use ('cuda', 'cpu', or None for auto)
- `verbose`: Print detailed extraction information
- `cache_model`: Cache model for faster repeated use

**Returns:**
- NumPy array of shape `(num_layers, hidden_size)` or `(hidden_size,)` for single layer

## 📊 Output Shape Reference

| Model | Hidden Size | Output Shape (3 layers) |
|-------|------------|-------------------------|
| GPT-2 | 768 | (3, 768) |
| GPT-2 Medium | 1024 | (3, 1024) |
| GPT-2 Large | 1280 | (3, 1280) |
| GPT-2 XL | 1600 | (3, 1600) |
| BERT Base | 768 | (3, 768) |
| BERT Large | 1024 | (3, 1024) |

## 🔧 Requirements

- Python 3.8+
- PyTorch
- Transformers
- NNsight
- NumPy
- inflect (for grammar utilities)

## 📈 Use Cases

- **Interpretability Research** - Analyze internal representations of LLMs
- **Probing Classifiers** - Extract features for downstream tasks
- **Semantic Similarity** - Compare representations across texts
- **Layer Analysis** - Study information flow through model layers
- **Neuron Analysis** - Investigate activation patterns
- **Model Debugging** - Understand model behavior on specific inputs

## 🤝 Contributing

Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.

## 📝 Citation

If you use `nnez` in your research, please cite:

```bibtex
@software{nnez2024,
  title = {nnez: Neural Network Easy Extraction},
  author = {Your Name},
  year = {2024},
  url = {https://github.com/yourusername/nnez}
}
```

## 📄 License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

## 🙏 Acknowledgments

- Built on top of [NNsight](https://github.com/ndif-team/nnsight) for model introspection
- Uses [inflect](https://github.com/jaraco/inflect) for grammar utilities
- Compatible with [HuggingFace Transformers](https://github.com/huggingface/transformers)

## 🔗 Links

- [PyPI Package](https://pypi.org/project/nnez/)
- [GitHub Repository](https://github.com/yourusername/nnez)
- [Documentation](https://nnez.readthedocs.io/)
- [Issue Tracker](https://github.com/yourusername/nnez/issues)

---

Made with ❤️ for the ML interpretability community
