$ show center-for-humans-and-machines/transformer-heads
The transformer_heads library is a Python toolkit for attaching auxiliary heads to transformer models, enabling tasks such as linear probing, fine-tuning, and multi-task learning. Its core functionality revolves around a modular design where users can define configurable heads (via HeadConfig) that are hooked into intermediate layers of a base transformer model. The library integrates with Hugging Face Transformers, PEFT (for LoRA/QLoRA), and bitsandbytes (for quantization) to support efficient adaptation and deployment. Architecturally, the library consists of several key components: a configuration system (config.py), flexible head implementations (model/head.py), a model wrapper that manages head attachment and forward hooks (model/model.py), and utilities for loading, evaluation, and data handling. Notable patterns include the use of PyTorch forward hooks to capture intermediate representations without modifying base weights, and the integration of PEFT methods with custom heads for parameter-efficient fine-tuning. The library also provides custom saving/loading mechanisms to persist both base model adaptations and head weights. While the library demonstrates strong software engineering practices and reusable patterns, it is currently deprecated and suffers from high technical drift. The README explicitly states that it is broken with newer versions of the Transformers library and is no longer maintained. This limits its usability in contemporary workflows, despite the soundness of its core concepts. The included Jupyter notebooks (e.g., for GPT-2 and Llama models) serve as educational examples but may require dependency adjustments to run successfully.
The LLM identified 3 files worth highlighting for reuse:
Defines reusable configuration classes (HeadConfig, HeadedConfig) for specifying head properties (layer, size, activation, loss function). Decoupled from model implementation, allowing easy extraction for other modular systems.
Implements MLPHead, a flexible multi-layer perceptron head with configurable layers, activation, and dropout. Self-contained and reusable as a building block for task-specific heads in any neural network architecture.
Implements TransformerWithHeads class that attaches auxiliary heads to intermediate transformer layers via forward hooks. Demonstrates a non-obvious pattern for modifying model behavior without altering base weights, enabling probing and multi-task learning.