Fela 1.6M

README.md

@@ -0,0 +1,49 @@
+---
+license: mit
+tags:
+- protein
+- biology
+- language-model
+- causal-lm
+- hyena
+- pytorch
+pipeline_tag: feature extraction
+---
+
+# Fela
+
+PyTorch written protein language model on the hyena operator (1.6M params)
+
+- Architecture: long conv + MLP blocks, pre-norm, LM head
+- Tokenizer: char level over `ACDEFGHIKLMNPQRSTVWYX`, `<pad>`=0, `<eos>`=22, `<unk>`=23
+- Data: Pfam-A (filtered to 20–512 residues, standard alphabet only), ~9.5B tokens
+- Training: 40k steps, batch 256, bf16, AdamW (wd 0.1), cosine LR 6e-4 → 6e-5
+
+## Config
+
+| Parameter | Value |
+|---|---|
+| d_model | 256 |
+| n_layer | 2 |
+| d_inner | 1024 |
+| vocab_size | 32 |
+| l_max | 514 |
+| order | 2 |
+| filter_order | 64 |
+| short_filter_order | 3 |
+| emb_dim | 5 |
+| w | 10 |
+| num_inner_mlps | 2 |
+| residual_in_fp32 | true |
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("pandeyps/fela", trust_remote_code=True)
+tok = AutoTokenizer.from_pretrained("pandeyps/fela", trust_remote_code=True)
+
+ids = tok.encode("MSDKIIEYDETARRAIEAGVNTLADAV", return_tensors="pt")
+gen = model.generate(ids, max_new_tokens=64, do_sample=True, temperature=0.7)
+print(tok.decode(gen[0]))

config.json

@@ -0,0 +1,20 @@
+{
+  "d_inner": 1024,
+  "d_model": 256,
+  "embed_dropout": 0.1,
+  "eos_token_id": 22,
+  "eps": 1e-05,
+  "l_max": 514,
+  "model_type": "fela",
+  "n_layer": 2,
+  "pad_token_id": 0,
+  "resid_dropout": 0.0,
+  "residual_in_fp32": true,
+  "transformers_version": "4.57.6",
+  "vocab_size": 32,
+  "auto_map": {
+    "AutoConfig": "modeling_fela.FelaConfig",
+    "AutoModelForCausalLM": "modeling_fela.FelaForCausalLM",
+    "AutoTokenizer": "tokenization_fela.FelaTokenizer"
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd95f80847e19fc12494f869447d77803204bb460421fb56def031bca5e59546
+size 6613656

modeling_fela.py

@@ -0,0 +1,224 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutput
+
+
+class Sin(nn.Module):
+    def __init__(self, dim, w=10, train_freq=True):
+        super().__init__()
+        self.freq = nn.Parameter(w * torch.ones(1, dim)) if train_freq else w * torch.ones(1, dim)
+    def forward(self, x):
+        return torch.sin(self.freq * x)
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, emb_dim, seq_len):
+        super().__init__()
+        t = torch.linspace(0, 1, seq_len)[None, :, None]
+        bands = (emb_dim - 1) // 2
+        t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]
+        w = 2 * math.pi * t_rescaled / seq_len
+        f = torch.linspace(1e-4, bands - 1, bands)[None, None]
+        z = torch.exp(-1j * f * w)
+        z = torch.cat([t, z.real, z.imag], dim=-1)
+        self.register_buffer("z", z)
+        self.register_buffer("t", t)
+    def forward(self, L):
+        return self.z[:, :L], self.t[:, :L]
+
+class ExponentialModulation(nn.Module):
+    def __init__(self, d_model, fast_decay_pct=0.3, slow_decay_pct=1.5, target=1e-2, shift=0.0):
+        super().__init__()
+        self.shift = shift
+        max_decay = math.log(target) / fast_decay_pct
+        min_decay = math.log(target) / slow_decay_pct
+        deltas = torch.linspace(min_decay, max_decay, d_model)[None, None]
+        self.register_buffer("deltas", deltas)
+    def forward(self, t, x):
+        return x * (torch.exp(-t * self.deltas.abs()) + self.shift)
+
+class HyenaFilter(nn.Module):
+    def __init__(self, d_model=256, emb_dim=5, order=64, seq_len=514,
+                 num_inner_mlps=2, w=10, modulate=True):
+        super().__init__()
+        self.modulate = modulate
+        self.bias = nn.Parameter(torch.randn(d_model))
+        act = Sin(dim=order, w=w)
+        self.pos_emb = PositionalEmbedding(emb_dim, seq_len)
+        self.implicit_filter = nn.Sequential(
+            nn.Linear(emb_dim, order), act,
+            *[m for _ in range(num_inner_mlps) for m in (nn.Linear(order, order), act)],
+            nn.Linear(order, d_model, bias=False),
+        )
+        self.modulation = ExponentialModulation(d_model)
+    def filter(self, L):
+        z, t = self.pos_emb(L)
+        h = self.implicit_filter(z)
+        if self.modulate:
+            h = self.modulation(t, h)
+        return h
+
+class ShortConv(nn.Module):
+    def __init__(self, d_model=256, order=2, short_filter_order=3):
+        super().__init__()
+        total_width = d_model * (order + 1)
+        self.in_proj = nn.Linear(d_model, total_width)
+        self.conv = nn.Conv1d(total_width, total_width, short_filter_order,
+                              groups=total_width, padding=short_filter_order - 1)
+    def forward(self, u):
+        u = self.in_proj(u)
+        u = u.transpose(1, 2)
+        u = self.conv(u)[..., :u.shape[-1]]
+        return u
+
+def fft_conv(u, k, bias=None):
+    seqlen = u.shape[-1]
+    fft_size = 2 * seqlen
+    k_f = torch.fft.rfft(k, n=fft_size) / fft_size
+    if len(u.shape) > 3:
+        k_f = k_f.unsqueeze(1)
+    u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)
+    y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen]
+    return y + u * bias.unsqueeze(-1) if bias is not None else y
+
+class HyenaOperator(nn.Module):
+    def __init__(self, d_model=256, l_max=514, order=2, filter_order=64,
+                 short_filter_order=3, drop_rate=0.0):
+        super().__init__()
+        self.d_model, self.order, self.l_max = d_model, order, l_max
+        self.in_proj = nn.Linear(d_model, (order + 1) * d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        total_width = d_model * (order + 1)
+        self.short_filter = nn.Conv1d(total_width, total_width, short_filter_order,
+                                      groups=total_width, padding=short_filter_order - 1)
+        self.filter_fn = HyenaFilter(d_model=d_model, order=filter_order, seq_len=l_max)
+        self.dropout = nn.Dropout(drop_rate)
+    def forward(self, u):
+        l_filter = min(u.size(-2), self.l_max)
+        u = rearrange(self.in_proj(u), "b l d -> b d l")
+        uc = self.short_filter(u)[..., :l_filter]
+        *x, v = uc.split(self.d_model, dim=1)
+        k = self.filter_fn.filter(l_filter)
+        k = rearrange(k, "c l (v o) -> c o v l", v=self.d_model, o=self.order - 1)
+        bias = rearrange(self.filter_fn.bias, "(v o) -> o v", o=self.order - 1)
+        for o, x_i in enumerate(reversed(x[1:])):
+            v = self.dropout(v * x_i)
+            v = fft_conv(v, k[o], bias[o])
+        return self.out_proj(rearrange(v * x[0], "b v l -> b l v"))
+
+class _Block(nn.Module):
+    def __init__(self, d_model, d_inner, l_max, drop1_p, drop2_p, eps=1e-5, residual_in_fp32=True):
+        super().__init__()
+        self.drop1 = nn.Dropout(drop1_p)
+        self.norm1 = nn.LayerNorm(d_model, eps=eps)
+        self.mixer = HyenaOperator(d_model=d_model, l_max=l_max)
+        self.drop2 = nn.Dropout(drop2_p)
+        self.norm2 = nn.LayerNorm(d_model, eps=eps)
+        self.mlp = nn.Sequential(nn.Linear(d_model, d_inner),
+                                 nn.GELU(approximate="tanh"),
+                                 nn.Linear(d_inner, d_model))
+        self.residual_in_fp32 = residual_in_fp32
+    def forward(self, hidden, residual):
+        dropped = self.drop1(hidden)
+        residual = (dropped + residual) if residual is not None else dropped
+        hidden = self.mixer(self.norm1(residual.to(dtype=self.norm1.weight.dtype)))
+        if self.residual_in_fp32:
+            residual = residual.float()
+        dropped = self.drop2(hidden)
+        residual = (dropped + residual) if residual is not None else dropped
+        hidden = self.mlp(self.norm2(residual.to(dtype=self.norm2.weight.dtype)))
+        if self.residual_in_fp32:
+            residual = residual.float()
+        return hidden, residual
+
+class Fela(nn.Module):
+    def __init__(self, d_model=256, n_layer=2, d_inner=1024, vocab_size=32, l_max=514,
+                 embed_dropout=0.1, resid_dropout=0.0, eps=1e-5, residual_in_fp32=True):
+        super().__init__()
+        torch.manual_seed(2222)
+        self.embed = nn.Embedding(vocab_size, d_model)
+        self.blocks = nn.ModuleList(
+            _Block(d_model, d_inner, l_max,
+                   drop1_p=embed_dropout if i == 0 else resid_dropout,
+                   drop2_p=resid_dropout, eps=eps,
+                   residual_in_fp32=residual_in_fp32)
+            for i in range(n_layer)
+        )
+        self.drop_f = nn.Dropout(resid_dropout)
+        self.ln_f = nn.LayerNorm(d_model, eps=eps)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        self._init_weights(n_layer)
+        self.lm_head.weight = self.embed.weight
+    def _init_weights(self, n_layer):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, std=0.02)
+        for name, p in self.named_parameters():
+            if name.endswith("out_proj.weight") or name.endswith("mlp.2.weight"):
+                nn.init.normal_(p, std=0.02 / math.sqrt(2 * n_layer))
+    def forward(self, input_ids):
+        hidden = self.embed(input_ids)
+        residual = None
+        for block in self.blocks:
+            hidden, residual = block(hidden, residual)
+        dropped = self.drop_f(hidden)
+        residual = (dropped + residual) if residual is not None else dropped
+        hidden = self.ln_f(residual.to(dtype=self.ln_f.weight.dtype))
+        return self.lm_head(hidden)
+
+
+try:
+    from transformers.generation import GenerationMixin
+except ImportError:
+    from transformers.generation_utils import GenerationMixin
+
+class FelaConfig(PretrainedConfig):
+    model_type = "fela"
+    def __init__(self, d_model=256, n_layer=2, d_inner=1024, vocab_size=32, l_max=514,
+                 embed_dropout=0.1, resid_dropout=0.0, eps=1e-5, residual_in_fp32=True, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model, self.n_layer, self.d_inner = d_model, n_layer, d_inner
+        self.vocab_size, self.l_max = vocab_size, l_max
+        self.embed_dropout, self.resid_dropout = embed_dropout, resid_dropout
+        self.eps, self.residual_in_fp32 = eps, residual_in_fp32
+        self.pad_token_id = 0
+        self.eos_token_id = 22
+
+class FelaPreTrainedModel(PreTrainedModel):
+    config_class = FelaConfig
+    base_model_prefix = "fela"
+    def _init_weights(self, module):
+        pass
+
+class FelaForCausalLM(FelaPreTrainedModel, GenerationMixin):
+    config_class = FelaConfig
+    base_model_prefix = "fela"
+    def __init__(self, config):
+        super().__init__(config)
+        self.fela = Fela(
+            d_model=config.d_model, n_layer=config.n_layer,
+            d_inner=config.d_inner, vocab_size=config.vocab_size,
+            l_max=config.l_max, embed_dropout=config.embed_dropout,
+            resid_dropout=config.resid_dropout, eps=config.eps,
+            residual_in_fp32=config.residual_in_fp32,
+        )
+    def get_input_embeddings(self):
+        return self.fela.embed
+    def set_input_embeddings(self, v):
+        self.fela.embed = v
+        self.fela.lm_head.weight = v.weight
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+    def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
+        logits = self.fela(input_ids)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.reshape(-1, self.config.vocab_size), labels.reshape(-1))
+        return CausalLMOutput(logits=logits, loss=loss)

special_tokens_map.json

@@ -0,0 +1 @@
+{}

tokenization_fela.py

@@ -0,0 +1,25 @@
+
+import json, os
+from transformers import PreTrainedTokenizer
+
+class FelaTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    def __init__(self, vocab_file=None, **kwargs):
+        self.vocab = json.load(open(vocab_file)) if vocab_file else {}
+        self.itos = {v: k for k, v in self.vocab.items()}
+        super().__init__(**kwargs)
+    def _tokenize(self, text, **kwargs):
+        return list(text)
+    def _convert_token_to_id(self, token):
+        return self.vocab.get(token, self.vocab.get("<unk>", 23))
+    def _convert_id_to_token(self, index):
+        return self.itos.get(index, "<unk>")
+    def get_vocab(self):
+        return dict(self.vocab)
+    def vocab_size(self):
+        return len(self.vocab)
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        fname = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json")
+        with open(fname, "w") as f:
+            json.dump(self.get_vocab(), f)
+        return (fname,)

tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "added_tokens_decoder": {},
+  "clean_up_tokenization_spaces": false,
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "FelaTokenizer"
+}

vocab.json

@@ -0,0 +1 @@
+{"A": 1, "C": 2, "D": 3, "E": 4, "F": 5, "G": 6, "H": 7, "I": 8, "K": 9, "L": 10, "M": 11, "N": 12, "P": 13, "Q": 14, "R": 15, "S": 16, "T": 17, "V": 18, "W": 19, "Y": 20, "X": 21, "<pad>": 0, "<eos>": 22, "<unk>": 23}