fast_cooc

GPU-first co-occurrence matrix builder for large text corpora. Uses CUDA kernels via CuPy to count word co-occurrences, with automatic dense/sparse mode selection based on vocabulary size and available VRAM.

Installation

pip install .

For HuggingFace tokenizer support:

pip install ".[hf]"

Requirements: CUDA-capable GPU, Python 3.10+, CUDA 12.x toolkit.

Quickstart

CLI

# Wikipedia with default settings (auto mode, regex tokenizer)
fast-cooc --dataset nlp/wikipedia --workdir wiki_out

# Small vocab → auto-selects dense mode (fast atomicAdd path)
fast-cooc --dataset nlp/wikipedia --max-vocab 30000 --workdir wiki_out

# Large vocab → auto-selects sparse mode
fast-cooc --dataset nlp/wikipedia --max-vocab 300000 --workdir wiki_out

# Force a specific mode
fast-cooc --dataset nlp/wikipedia --mode dense --max-vocab 30000

# Use a HuggingFace tokenizer
fast-cooc --dataset nlp/wikipedia --tokenizer hf --hf-model bert-base-uncased

Python API

from fast_cooc import run_warp_pipeline, VocabConfig, CountConfig

out_path = run_warp_pipeline(
    dataset_id="nlp/wikipedia",
    workdir="wiki_out",
    row_limit=100_000,                       # None for full dataset
    vocab_cfg=VocabConfig(min_freq=10, max_vocab=30_000),
    count_cfg=CountConfig(window=5, max_vram_gb=20.0, mode="auto"),
)
# out_path is either wiki_out/embeddings.npy (dense) or wiki_out/cooc_final.npz (sparse)

Using the GPU kernel directly

import numpy as np
from fast_cooc import count_cooc_gpu, collapse

# corpus_ids: int32 array where each value is a word index (or -1 for OOV)
corpus_ids = np.array([0, 1, 2, 3, 1, 0, 2, 3, 1, 2], dtype=np.int32)
n_vocab = 4

# Dense mode — returns a CuPy (n_vocab, n_vocab) float32 array on GPU
matrix = count_cooc_gpu(corpus_ids, n_vocab, window=2, mode="dense")

# Collapse: log1p → center columns → L2 normalize rows → CPU numpy array
embeddings = collapse(matrix)

# Sparse mode — returns a scipy.sparse.csr_matrix
sparse_matrix = count_cooc_gpu(corpus_ids, n_vocab, window=2, mode="sparse")

Custom tokenizers

from fast_cooc import run_warp_pipeline, RegexTokenizer, HFTokenizer, make_tokenizer

# Regex (default)
tok = RegexTokenizer(pattern=r"[A-Za-z0-9_']+", lowercase=True)

# HuggingFace subword tokenizer
tok = HFTokenizer(model_name="bert-base-uncased")

# Factory function
tok = make_tokenizer("regex")
tok = make_tokenizer("hf", hf_model="bert-base-uncased")
tok = make_tokenizer("callable", fn=lambda text: text.lower().split())

out = run_warp_pipeline("nlp/wikipedia", "wiki_out", tokenizer=tok)

Dense vs Sparse mode

mode="auto" (default) picks dense when the V x V matrix fits in 80% of the VRAM budget, sparse otherwise.

CLI reference

Pipeline stages

1. Build vocabulary     — tokenize dataset, count frequencies, filter by min_freq/max_vocab
2. Write ID memmap      — second pass: map tokens to int32 indices, write to disk
3. GPU co-occurrence    — chunk corpus, launch CUDA kernels, accumulate counts
4. Collapse (dense)     — log1p → column-center → L2 normalize → embeddings.npy
   Save (sparse)        — raw counts → cooc_final.npz

Output files

All outputs are written to --workdir:

API reference

run_warp_pipeline

run_warp_pipeline(
    dataset_id: str,
    workdir: str,
    row_limit: Optional[int] = None,
    vocab_cfg: VocabConfig = VocabConfig(),
    count_cfg: CountConfig = CountConfig(),
    tokenizer: Optional[TokenizerBackend] = None,
) -> Path

End-to-end pipeline. Returns path to the output file.

count_cooc_gpu

count_cooc_gpu(
    corpus_ids: np.ndarray,           # int32 token IDs (-1 = OOV)
    n_vocab: int,
    window: int = 5,
    max_vram_gb: float = 20.0,
    mode: Literal["auto", "dense", "sparse"] = "auto",
    flush_triplets_every: int = 20_000_000,
    progress_every_tokens: int = 0,
) -> Union[cp.ndarray, sp.csr_matrix]

GPU co-occurrence counting. Returns CuPy dense array or scipy sparse CSR matrix.

collapse

collapse(matrix: cp.ndarray) -> np.ndarray

Transforms a dense co-occurrence matrix into normalized embeddings on GPU: log1p → subtract column means → L2 normalize rows. Returns a CPU numpy array.

Config

VocabConfig(min_freq: int = 10, max_vocab: int = 300_000)
CountConfig(window: int = 5, max_vram_gb: float = 20.0, mode: str = "auto", flush_triplets_every: int = 20_000_000)

Tokenizers

All tokenizers implement the TokenizerBackend protocol (tokenize(text: str) -> Iterable[str]).