File size: 3,677 Bytes
36b95ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Config for the Fly connectome language models (published with trust_remote_code)."""
from transformers import PretrainedConfig
CONNECTOME = ("MaleCNS v1.0, FlyEM / HHMI Janelia, University of Cambridge, MRC LMB, "
"Google Research. CC BY 4.0.")
class FlyConfig(PretrainedConfig):
"""A fruit-fly connectome as the recurrent layer of a language model.
Two architectures share this config:
* ``decoder`` — token in, token out, nothing else. Comparable to ngxson/fly-llm-hf.
* ``encoder_decoder`` — the post being answered is encoded once by a frozen sentence encoder
(``news_encoder``) and injected as a constant current into the olfactory neurons for the whole
generation, so the fly "smells" the post while it writes. Only the projection into those
neurons is trained here; the sentence encoder is a separate model and is not shipped.
The recurrent weights are anatomy, not learned: ``w_values`` are signed, row-normalised synapse
counts between the neurons of the release above. ``mode`` records whether the published values
are the raw connectome (``gains``) or connectome-masked learned strengths (``edges``).
"""
model_type = "fly"
keys_to_ignore_at_inference = ["cache_params"]
def __init__(
self,
arch="decoder",
vocab_size=2048,
d_emb=256,
delay=8,
ticks=2,
mode="gains",
token_input="cb_sensory",
n_neurons=1,
n_edges=0,
n_token_input=1,
n_news_input=0,
readout="all",
readout_size=1,
readout_rank=256,
head_type="lowrank",
news_dim=0,
news_group="orn",
n_reserved=0,
news_mode="direct",
news_glom=0,
news_encoder="none",
news_prefix="",
min_syn=1,
modulatory_sign=0.0,
connectome=CONNECTOME,
source_repo="",
use_cache=True,
pad_token_id=0,
bos_token_id=1,
eos_token_id=2,
tie_word_embeddings=False,
**kwargs,
):
self.arch = arch
self.vocab_size = vocab_size
self.d_emb = d_emb
self.delay = delay
self.ticks = ticks
self.mode = mode
self.token_input = token_input # which population the delay line writes into
self.n_neurons = n_neurons
self.n_edges = n_edges
self.n_token_input = n_token_input
self.n_news_input = n_news_input
self.readout = readout
self.readout_size = readout_size
self.readout_rank = readout_rank
self.head_type = head_type # lowrank: Linear->LayerNorm->Linear, linear: LayerNorm->Linear
self.news_dim = news_dim
self.news_group = news_group # population carved out for the post channel
self.n_reserved = n_reserved # its size, even when this variant does not use it
self.news_mode = news_mode
self.news_glom = news_glom
self.news_encoder = news_encoder
self.news_prefix = news_prefix
self.min_syn = min_syn
self.modulatory_sign = modulatory_sign
self.connectome = connectome
self.source_repo = source_repo
self.use_cache = use_cache
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_word_embeddings,
**kwargs,
)
@property
def hidden_size(self):
return self.readout_size
@property
def conditional(self):
return self.arch == "encoder_decoder" and self.news_dim > 0
|