StructLens

Visualizer for StructLens: https://naist-nlp.github.io/structlens/

Installation

# uv
uv add git+https://github.com/naist-nlp/structlens.git
# pip
# pip install git+https://github.com/naist-nlp/structlens.git

Usage

import torch
from structlens import L2DistanceSimilarityFunction, StructLens, create_masks

num_tokens_per_layer = 5
num_layers = 3
hidden_size = 128

representations = torch.randn(num_layers, num_tokens_per_layer, hidden_size)

struct_lens = StructLens()
st_list = struct_lens(representations)
for i, st in enumerate(st_list):
    print("layer: ", i)
    print("argmax_heads: ", st.argmax_heads)

See examples/structlens-sample for more details.

Development

Setup

Install dependencies:

uv sync

Prepare pre-commit:

uv run lefthook install

Running Tests

uv run pytest

Formatting & Linting

uv run ruff check --fix

Type Checking

uv run ty check