Jathurshan/TFM-Tokenizer

Tokenizing Single-Channel EEG with Time-Frequency Motif Learning - ICLR 2026

Official pretrained and finetuned weights for Tokenizing Single-Channel EEG with Time-Frequency Motif Learning, published at ICLR 2026.

Repository Structure

pretrained/
├── tfm_tokenizer_last.pth            # TFM-Tokenizer (VQ-VAE, 2×2×8)
└── tfm_encoder_mtp_last.pth          # TFM-Encoder pretrained via MTP

finetuned/
├── TUEV/seed_{1..5}/best_model.pth   # 6-class EEG event detection
├── TUAB/seed_{1..5}/best_model.pth   # Binary abnormal EEG detection
└── CHBMIT/seed_{1..5}/best_model.pth # Binary seizure detection

models/
└── tfm_token.py                      # Model definitions

---

Quick start

Installation

pip install torch einops linear_attention_transformer huggingface_hub

Clone the source repository for model definitions and utilities:

git clone https://github.com/Jathurshan0330/TFM-Tokenizer.git
cd TFM-Tokenizer

---

1. Load the pretrained TFM-Tokenizer

import torch
from huggingface_hub import hf_hub_download
from models.tfm_token import get_tfm_tokenizer_2x2x8
from utils.utils import get_stft_torch

ckpt = hf_hub_download(repo_id="Jathurshan/TFM-Tokenizer", filename="pretrained/tfm_tokenizer_last.pth")
tokenizer = get_tfm_tokenizer_2x2x8(code_book_size=8192, emb_size=64)
tokenizer.load_state_dict(torch.load(ckpt, map_location="cpu"))
tokenizer.eval()

2. Load the MTP-pretrained TFM-Encoder (for finetuning on a new task)

from models.tfm_token import get_tfm_token_classifier_64x4

ckpt = hf_hub_download(repo_id="Jathurshan/TFM-Tokenizer", filename="pretrained/tfm_encoder_mtp_last.pth")
model = get_tfm_token_classifier_64x4(n_classes=YOUR_NUM_CLASSES, code_book_size=8192, emb_size=64)

checkpoint = torch.load(ckpt, map_location="cpu")
filtered = {k: v for k, v in checkpoint.items() if "classification_head" not in k}
model.load_state_dict(filtered, strict=False)
# classification_head is randomly initialized — finetune on your data

3. Load a finetuned checkpoint (for direct inference)

# Example: TUEV dataset, seed 1
ckpt = hf_hub_download(repo_id="Jathurshan/TFM-Tokenizer", filename="finetuned/TUEV/seed_1/best_model.pth")
model = get_tfm_token_classifier_64x4(n_classes=6, code_book_size=8192, emb_size=64)
model.load_state_dict(torch.load(ckpt, map_location="cpu"))
model.eval()

Dataset-specific n_classes:

• TUEV: n_classes=6 (multi-class)
• TUAB: n_classes=1 (binary, use sigmoid)
• CHBMIT: n_classes=1 (binary, use sigmoid)

4. Full inference pipeline

import torch
from einops import rearrange
from huggingface_hub import hf_hub_download
from models.tfm_token import get_tfm_tokenizer_2x2x8, get_tfm_token_classifier_64x4
from utils.utils import get_stft_torch

# Load tokenizer
tok_ckpt = hf_hub_download(repo_id="Jathurshan/TFM-Tokenizer", filename="pretrained/tfm_tokenizer_last.pth")
tokenizer = get_tfm_tokenizer_2x2x8(code_book_size=8192, emb_size=64)
tokenizer.load_state_dict(torch.load(tok_ckpt, map_location="cpu"))
tokenizer.eval()

# Load finetuned encoder (e.g. TUEV seed 1)
enc_ckpt = hf_hub_download(repo_id="Jathurshan/TFM-Tokenizer", filename="finetuned/TUEV/seed_1/best_model.pth")
encoder = get_tfm_token_classifier_64x4(n_classes=6, code_book_size=8192, emb_size=64)
encoder.load_state_dict(torch.load(enc_ckpt, map_location="cpu"))
encoder.eval()

# Inference on raw EEG: x shape (B, C, T) at 200 Hz
x_temporal = x
B, C, T = x_temporal.shape
x_stft = get_stft_torch(x_temporal, resampling_rate=200)
x_stft = rearrange(x_stft, 'B C F T -> (B C) F T')
x_temporal_flat = rearrange(x_temporal, 'B C T -> (B C) T')

with torch.no_grad():
    _, x_tokens, _ = tokenizer.tokenize(x_stft, x_temporal_flat)
    x_tokens = rearrange(x_tokens, '(B C) T -> B C T', C=C)
    preds = encoder(x_tokens, num_ch=C)

---

Citation

If you find this work useful, please cite:

@inproceedings{
pradeepkumar2026tokenizing,
title={Tokenizing Single-Channel {EEG} with Time-Frequency Motif Learning},
author={Jathurshan Pradeepkumar and Xihao Piao and Zheng Chen and Jimeng Sun},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=2sPmWHZ8Ir}
}