untrained initialization, seed 1

init/README.md

@@ -0,0 +1,158 @@
+---
+license: cc-by-4.0
+language: [en]
+library_name: transformers
+pipeline_tag: text-generation
+base_model: QuixiAI/MaleCNS
+base_model_relation: finetune
+datasets: [karpathy/tiny_shakespeare]
+tags: [connectome, fruit-fly, drosophila, malecns, recurrent, sparse, tiny-shakespeare, custom_code]
+---
+
+# QuixiAI/FlyGPT
+
+A character-level language model whose recurrent architecture is **a real subgraph of the fruit-fly brain
+connectome** ([MaleCNS v1.0](https://male-cns.janelia.org/)). Unlike the earlier frozen-reservoir approach in
+[ngxson/fly-llm-hf](https://huggingface.co/ngxson/fly-llm-hf), which keeps the connectome's synaptic weights fixed and
+trains only the projections and readout, FlyGPT **trains one value per real synaptic connection with gradient
+descent** while keeping the fly's edge topology fixed, and compares the result against the same neurons with
+degree-preserving scrambled connections across paired seeds.
+
+**Base model: [QuixiAI/MaleCNS](https://huggingface.co/QuixiAI/MaleCNS)**, the lossless packaging of the MaleCNS v1.0
+connectivity tables. FlyGPT's graph is extracted from it deterministically (`build_graph.py`, revision pinned in
+`data/fly/build_edges.py`); `graph.node_id` and `graph.synapse_count` map every edge back to that repository.
+
+This checkpoint's wiring is the original MaleCNS wiring.
+
+**Untrained initialization.** The recurrent edge values and adapters are the exact random initialization that FlyGPT's `train.py` starts from for seed 1. Text generated from this checkpoint is noise by design.
+
+This is not a biological simulation of a living fly. The "weights" in the MaleCNS release are anatomical synapse
+counts; they are stored here as `graph.synapse_count` and are **not** the model's parameters.
+
+## The graph
+
+Every number below is produced by FlyGPT's extraction script (`build_graph.py`), not typed by hand.
+
+| | |
+|---|---|
+| Source | MaleCNS v1.0 flat connectome (`gs://flyem-male-cns/v1.0/connectome-data/flat-connectome/`) |
+| Candidate pool | central brain: `superclass` starting with `cb_` (37,108 neurons) |
+| Minimum synapses per connection | 3 (engineering choice, not a biological claim) |
+| Extraction | largest SCC → largest directed (k,k)-core with ≥ target nodes (k = 40) → trim by weighted degree |
+| Neurons used | 5,000 |
+| Directed connections used | 524,324 |
+| Synaptic contacts represented | 8,300,915 |
+| Largest SCC fraction | 1.0 |
+| Reciprocal pairs | 93,055 |
+| Input / output neurons | top 256 by out-degree / top 512 by in-degree |
+| Input→output shortest path (median / p90 / max hops) | 1.0 / 1.0 / 1.0 |
+| Graph hash | `f82b783b7ccb5a354fc4cf3de6de4a98d75029303c55f8faae28ab807828a007` |
+
+`graph.node_id` holds the MaleCNS body ids, so every neuron maps back to the release.
+
+## What is in `model.safetensors`
+
+| tensor | shape | dtype | size |
+|---|---|---|---|
+| `graph.edge_index` | (2, 524324) | int32 | 4.19 MB |
+| `graph.synapse_count` | (524324,) | int32 | 2.10 MB |
+| `graph.node_id` | (5000,) | int64 | 0.04 MB |
+| `graph.input_nodes` | (256,) | int64 | 0.00 MB |
+| `graph.output_nodes` | (512,) | int64 | 0.00 MB |
+| `recurrent.edge_values` | (524324,) | bfloat16 | 1.05 MB |
+| `recurrent.bias` | (5000,) | bfloat16 | 0.01 MB |
+| `recurrent.raw_leak` | (5000,) | bfloat16 | 0.01 MB |
+| `embed.weight` | (65, 32) | bfloat16 | 0.00 MB |
+| `input_proj.weight` | (256, 32) | bfloat16 | 0.02 MB |
+| `input_proj.bias` | (256,) | bfloat16 | 0.00 MB |
+| `lm_head.weight` | (65, 512) | bfloat16 | 0.07 MB |
+| `lm_head.bias` | (65,) | bfloat16 | 0.00 MB |
+
+`graph.*` is the anatomy (integer, never trained). `recurrent.*`, `embed.*`, `input_proj.*`, `lm_head.*` are the
+learned state, stored in bf16. The sparse recurrent matmul is rebuilt in fp32 at runtime (rows = destination,
+columns = source), with each incoming edge scaled by `1/sqrt(in_degree)`.
+
+## Dynamics
+
+```text
+character → embedding (32) → linear → 256 input neurons
+proposal_i = tanh( Σ_j W_ij h_j / sqrt(in_degree_i) + external_input_i + bias_i )
+h_i ← (1 − leak_i) h_i + leak_i · proposal_i        (2 microsteps per character, leak_i = sigmoid(raw_leak_i))
+512 output neuron states → linear → 65 logits
+```
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tok = AutoTokenizer.from_pretrained("QuixiAI/FlyGPT")  # tokenizer is shared
+model = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", subfolder="init", trust_remote_code=True, dtype=torch.float32)
+
+ids = tok("ROMEO:", return_tensors="pt").input_ids
+out = model.generate(ids, max_new_tokens=300, do_sample=True, temperature=0.8)
+print(tok.decode(out[0]))
+```
+
+The tokenizer is strict: only the 65 characters of Tiny Shakespeare are encodable.
+
+## Citation
+
+If you use this model, please cite it, its base model, and the MaleCNS dataset paper.
+
+This model:
+
+```bibtex
+@misc{hartford2026flygpt,
+  title        = {FlyGPT: a language model whose recurrent architecture is a real subgraph of the fruit-fly connectome},
+  author       = {Hartford, Eric},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/QuixiAI/FlyGPT}},
+  note         = {Base model: QuixiAI/MaleCNS (MaleCNS v1.0, Berg et al. 2026, CC-BY 4.0). Code: https://github.com/QuixiAI/FlyGPT}
+}
+```
+
+The base model (lossless connectome packaging):
+
+```bibtex
+@misc{hartford2026malecns,
+  title        = {QuixiAI/MaleCNS: the MaleCNS v1.0 fruit-fly connectome as lossless Safetensors},
+  author       = {Hartford, Eric},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  doi          = {10.57967/hf/10410},
+  howpublished = {\url{https://huggingface.co/QuixiAI/MaleCNS}},
+  note         = {Repackaging of Berg et al. (2026), CC-BY 4.0}
+}
+```
+
+The dataset (required by the CC-BY 4.0 license):
+
+```bibtex
+@article{berg2026malecns,
+  title     = {Sexual dimorphism in the complete {Drosophila} male central nervous system connectome},
+  author    = {Berg, Stuart and Beckett, Isabella R. and Costa, Marta and Schlegel, Philipp and Januszewski, Michał and Marin, Elizabeth C. and Nern, Aljoscha and Preibisch, Stephan and Qiu, Wei and Takemura, Shin-ya and Fragniere, Alexandra M.C. and Champion, Andrew S. and Adjavon, Diane-Yayra and Cook, Michael and Gkantia, Marina and Hayworth, Kenneth J. and Huang, Gary B. and Katz, William T. and Kämpf, Florian and Lu, Zhiyuan and Ordish, Christopher and Paterson, Tyler and Stürner, Tomke and Trautman, Eric T. and Whittle, Catherine R. and Burnett, Laura E. and Hoeller, Judith and Li, Feng and Loesche, Frank and Morris, Billy J. and Pietzsch, Tobias and Pleijzier, Markus W. and Silva, Valeria and Yin, Yijie and Ali, Iris and Badalamente, Griffin and Bates, Alexander Shakeel and Beresford, Rory J. and Bogovic, John and Brooks, Paul and Cachero, Sebastian and Canino, Brandon S. and Chaisrisawatsuk, Bhumpanya and Clements, Jody and Crowe, Arthur and de Haan Vicente, Inês and Dempsey, Georgia and Donà, Erika and Dos Santos, Márcia and Dreher, Marisa and Dunne, Christopher R. and Eichler, Katharina and Finley-May, Samantha and Flynn, Miriam A. and Hameed, Imran and Hopkins, Gary Patrick and Hubbard, Philip M. and Kiassat, Ladann and Kovalyak, Julie and Lauchie, Shirley A. and Leonard, Meghan and Lohff, Alanna and Longden, Kit D. and Maldonado, Charli A. and Moitra, Ilina and Moon, Sung Soo and Mooney, Caroline and Munnelly, Eva J. and Okeoma, Nneoma and Olbris, Donald J. and Pai, Anika and Patel, Birava and Phillips, Emily M. and Plaza, Stephen M. and Richards, Alana and Rivas Salinas, Jennifer and Roberts, Ruairí J.V. and Rogers, Edward M. and Scott, Ashley L. and Scuderi, Louis A. and Seenivasan, Pavithraa and Serratosa Capdevila, Laia and Smith, Claire and Svirskas, Rob and Takemura, Satoko and Tastekin, Ibrahim and Thomson, Alexander and Umayam, Lowell and Walsh, John J. and Whittome, Holly and Xu, C. Shan and Yakal, Emily A. and Yang, Tansy and Zhao, Arthur and George, Reed and Jain, Viren and Jayaraman, Vivek and Korff, Wyatt and Meissner, Geoffrey W. and Romani, Sandro and Funke, Jan and Knecht, Christopher and Saalfeld, Stephan and Scheffer, Louis K. and Waddell, Scott and Card, Gwyneth M. and Ribeiro, Carlos and Reiser, Michael B. and Hess, Harald F. and Rubin, Gerald M. and Jefferis, Gregory S.X.E.},
+  journal   = {Cell},
+  volume    = {189},
+  number    = {18},
+  pages     = {5504--5526.e15},
+  year      = {2026},
+  month     = sep,
+  publisher = {Elsevier},
+  doi       = {10.1016/j.cell.2026.08.015},
+  url       = {https://doi.org/10.1016/j.cell.2026.08.015},
+  note      = {Preprint: bioRxiv 10.1101/2025.10.09.680999. Data: MaleCNS v1.0, CC-BY 4.0, https://male-cns.janelia.org}
+}
+```
+
+## License
+
+The connectome is released under CC-BY 4.0 by the FlyEM Project Team (HHMI Janelia), the University of Cambridge,
+the MRC Laboratory of Molecular Biology, and Google Research. This checkpoint is a derivative and carries the same
+license.
+
+Prior art: [ngxson/fly-llm-hf](https://huggingface.co/ngxson/fly-llm-hf) (frozen MaleCNS reservoir LM) and
+[eob/gpt-fly](https://huggingface.co/eob/gpt-fly) (FlyWire-masked GPT-2). Code and experiment:
+[github.com/QuixiAI/FlyGPT](https://github.com/QuixiAI/FlyGPT).

init/config.json

@@ -0,0 +1,38 @@
+{
+  "model_type": "flygpt",
+  "architectures": [
+    "FlyGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_flygpt.FlyGPTConfig",
+    "AutoModelForCausalLM": "modeling_flygpt.FlyGPTForCausalLM"
+  },
+  "vocab_size": 65,
+  "num_neurons": 5000,
+  "num_edges": 524324,
+  "embedding_dim": 32,
+  "num_input_neurons": 256,
+  "num_output_neurons": 512,
+  "microsteps": 2,
+  "activation": "tanh",
+  "learned_leak": true,
+  "leak_init": 0.5,
+  "degree_normalization": true,
+  "init_scale": 1.0,
+  "dtype": "bfloat16",
+  "graph": {
+    "graph_name": "cb5k",
+    "condition": "real",
+    "control_seed": 1,
+    "source": "malecns-v1.0",
+    "region_filter": "central_brain",
+    "min_synapses": 3,
+    "hash": "f82b783b7ccb5a354fc4cf3de6de4a98d75029303c55f8faae28ab807828a007"
+  },
+  "training_state": {
+    "status": "init",
+    "condition": "real",
+    "seed": 1,
+    "project": "flygpt-v0"
+  }
+}

init/configuration_flygpt.py

@@ -0,0 +1,40 @@
+"""FlyGPT configuration for Hugging Face `transformers` (loaded with trust_remote_code=True)."""
+from transformers import PretrainedConfig
+
+
+class FlyGPTConfig(PretrainedConfig):
+    model_type = "flygpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 65,
+        num_neurons: int = 0,
+        num_edges: int = 0,
+        embedding_dim: int = 32,
+        num_input_neurons: int = 256,
+        num_output_neurons: int = 512,
+        microsteps: int = 2,
+        activation: str = "tanh",
+        learned_leak: bool = True,
+        leak_init: float = 0.5,
+        degree_normalization: bool = True,
+        init_scale: float = 0.1,
+        graph: dict | None = None,
+        training_state: dict | None = None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.num_neurons = num_neurons
+        self.num_edges = num_edges
+        self.embedding_dim = embedding_dim
+        self.num_input_neurons = num_input_neurons
+        self.num_output_neurons = num_output_neurons
+        self.microsteps = microsteps
+        self.activation = activation
+        self.learned_leak = learned_leak
+        self.leak_init = leak_init
+        self.degree_normalization = degree_normalization
+        self.init_scale = init_scale
+        self.graph = graph or {}                    # source, hash, region filter, stats, provenance
+        self.training_state = training_state or {}  # "init" or the run that produced the values
+        super().__init__(**kwargs)

init/generation_config.json

@@ -0,0 +1,5 @@
+{
+  "do_sample": true,
+  "temperature": 0.8,
+  "max_new_tokens": 300
+}

init/graph_metadata.json

@@ -0,0 +1,75 @@
+{
+  "graph_name": "cb5k",
+  "condition": "real",
+  "control_seed": 1,
+  "source": "malecns-v1.0",
+  "region_filter": "central_brain",
+  "min_synapses": 3,
+  "selector": "directed_core",
+  "target_neurons": 5000,
+  "hash": "f82b783b7ccb5a354fc4cf3de6de4a98d75029303c55f8faae28ab807828a007",
+  "stats": {
+    "candidate_pool": 37108,
+    "candidate_edges": 7294218,
+    "min_synapses": 3,
+    "target_neurons": 5000,
+    "edges_after_threshold": 2805370,
+    "scc_after_threshold": 36461,
+    "core_size_before_trim": 6576,
+    "trim_iterations": 1,
+    "core_k": 40,
+    "final_neurons": 5000,
+    "final_edges": 524324,
+    "source": "malecns-v1.0",
+    "n_neurons": 5000,
+    "n_edges": 524324,
+    "synaptic_contacts": 8300915.0,
+    "largest_scc_fraction": 1.0,
+    "reciprocal_pairs": 93055,
+    "in_degree": {
+      "mean": 104.8648,
+      "min": 4,
+      "max": 881
+    },
+    "out_degree": {
+      "mean": 104.8648,
+      "min": 5,
+      "max": 1539
+    },
+    "density": 0.020977155431086216,
+    "hash": "f82b783b7ccb5a354fc4cf3de6de4a98d75029303c55f8faae28ab807828a007",
+    "region_breakdown": {
+      "central_brain": 5000
+    }
+  },
+  "diagnostics": {
+    "provenance": {
+      "control": "real"
+    },
+    "io": {
+      "n_input": 256,
+      "n_output": 512,
+      "input_rule": "top_out_degree",
+      "output_rule": "top_in_degree",
+      "overlap": 160
+    },
+    "diagnostics": {
+      "n_neurons": 5000,
+      "n_edges": 524324,
+      "largest_scc_fraction": 1.0,
+      "reciprocal_pairs": 93055,
+      "io_reachable_fraction": 1.0,
+      "io_path_median": 1.0,
+      "io_path_p90": 1.0,
+      "io_path_max": 1.0
+    },
+    "gate": {
+      "microsteps": 2,
+      "max_p90_hops": 6,
+      "min_reachable_fraction": 0.95,
+      "reachable_ok": true,
+      "p90_ok": true,
+      "passed": true
+    }
+  }
+}

init/model.safetensors

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

init/modeling_flygpt.py

@@ -0,0 +1,164 @@
+"""FlyGPT: a character-level language model whose recurrent core is a real subgraph of the
+fruit-fly connectome (MaleCNS v1.0). Hugging Face `transformers` implementation; self-contained.
+
+Dynamics (one scalar state per neuron, plan.md §7 of the FlyGPT spec):
+
+    proposal_i = tanh( sum_j W_ij h_j / sqrt(in_degree_i) + external_input_i + bias_i )
+    h_i_new    = (1 - leak_i) * h_i + leak_i * proposal_i
+
+The connectome is stored in `model.safetensors` as integer tensors (`graph.*`); only the learned
+per-edge values and the adapters are floating point (bf16 on disk). The sparse recurrent matmul is
+rebuilt in fp32 at runtime (rows = destination, columns = source).
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.utils import ModelOutput
+
+from .configuration_flygpt import FlyGPTConfig
+
+
+@dataclass
+class FlyGPTOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    state: Optional[torch.FloatTensor] = None   # [B, N] neuron states after the last character
+
+
+class FlyGraph(nn.Module):
+    """The anatomy. Integer buffers only; never trained."""
+
+    def __init__(self, num_neurons: int, num_edges: int, num_input: int, num_output: int):
+        super().__init__()
+        self.register_buffer("edge_index", torch.zeros(2, num_edges, dtype=torch.int32))   # [source, destination]
+        self.register_buffer("synapse_count", torch.zeros(num_edges, dtype=torch.int32))   # MaleCNS synaptic contacts
+        self.register_buffer("node_id", torch.zeros(num_neurons, dtype=torch.int64))       # MaleCNS body ids
+        self.register_buffer("input_nodes", torch.zeros(num_input, dtype=torch.int64))
+        self.register_buffer("output_nodes", torch.zeros(num_output, dtype=torch.int64))
+
+
+class FlyRecurrentCore(nn.Module):
+    """The learned state: one value per real edge, plus per-neuron bias and leak."""
+
+    def __init__(self, num_neurons: int, num_edges: int, leak_init: float, learned_leak: bool):
+        super().__init__()
+        self.edge_values = nn.Parameter(torch.zeros(num_edges))
+        self.bias = nn.Parameter(torch.zeros(num_neurons))
+        self.raw_leak = nn.Parameter(torch.full((num_neurons,), math.log(leak_init / (1 - leak_init))),
+                                     requires_grad=learned_leak)
+
+
+class FlyGPTPreTrainedModel(PreTrainedModel):
+    config_class = FlyGPTConfig
+    base_model_prefix = "flygpt"
+    _is_stateful = True
+    _supports_cache_class = False
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, FlyRecurrentCore):
+            nn.init.normal_(module.edge_values, std=self.config.init_scale)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=1.0)
+
+
+class FlyGPTForCausalLM(FlyGPTPreTrainedModel, GenerationMixin):
+    def __init__(self, config: FlyGPTConfig):
+        super().__init__(config)
+        c = config
+        self.graph = FlyGraph(c.num_neurons, c.num_edges, c.num_input_neurons, c.num_output_neurons)
+        self.recurrent = FlyRecurrentCore(c.num_neurons, c.num_edges, c.leak_init, c.learned_leak)
+        self.embed = nn.Embedding(c.vocab_size, c.embedding_dim)
+        self.input_proj = nn.Linear(c.embedding_dim, c.num_input_neurons)
+        self.lm_head = nn.Linear(c.num_output_neurons, c.vocab_size)
+        self.post_init()
+
+    # ---- sparse recurrent matrix -------------------------------------------------------------
+    @property
+    def num_neurons(self) -> int:
+        return self.config.num_neurons
+
+    def edge_scale(self) -> torch.Tensor:
+        """1/sqrt(in_degree) per edge (degree normalization), or ones."""
+        dst = self.graph.edge_index[1].long()
+        if not self.config.degree_normalization:
+            return torch.ones_like(dst, dtype=torch.float32)
+        in_deg = torch.bincount(dst, minlength=self.num_neurons).clamp(min=1).float()
+        return 1.0 / in_deg[dst].sqrt()
+
+    def sparse_weight(self) -> torch.Tensor:
+        src, dst = self.graph.edge_index[0].long(), self.graph.edge_index[1].long()
+        values = self.recurrent.edge_values.float() * self.edge_scale()
+        return torch.sparse_coo_tensor(torch.stack([dst, src]), values, (self.num_neurons, self.num_neurons))
+
+    def dense_weight(self) -> torch.Tensor:
+        """Convenience for analysis; [N, N] with rows = destination. Never used in the forward pass."""
+        return self.sparse_weight().to_dense()
+
+    @property
+    def leak(self) -> torch.Tensor:
+        return torch.sigmoid(self.recurrent.raw_leak.float())
+
+    # ---- dynamics ------------------------------------------------------------------------------
+    def init_state(self, batch: int, device=None) -> torch.Tensor:
+        return torch.zeros(batch, self.num_neurons, device=device or self.recurrent.edge_values.device)
+
+    def drive(self, x: torch.Tensor) -> torch.Tensor:
+        d = torch.zeros(x.shape[0], self.num_neurons, device=x.device, dtype=torch.float32)
+        d[:, self.graph.input_nodes] = self.input_proj(self.embed(x)).float()
+        return d
+
+    def step(self, state: torch.Tensor, x: torch.Tensor, W: torch.Tensor | None = None) -> torch.Tensor:
+        W = self.sparse_weight() if W is None else W
+        drive, leak, bias = self.drive(x), self.leak, self.recurrent.bias.float()
+        for _ in range(self.config.microsteps):
+            incoming = torch.sparse.mm(W, state.float().T).T
+            proposal = torch.tanh(incoming + drive + bias)
+            state = (1 - leak) * state + leak * proposal
+        return state
+
+    def logits_from_state(self, state: torch.Tensor) -> torch.Tensor:
+        return self.lm_head(state[:, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float()
+
+    def forward(self, input_ids: torch.LongTensor, state: Optional[torch.Tensor] = None,
+                labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None,
+                return_dict: Optional[bool] = None, **kwargs) -> FlyGPTOutput:
+        B, T = input_ids.shape
+        state = self.init_state(B, input_ids.device) if state is None else state
+        W = self.sparse_weight()
+        outs = []
+        for t in range(T):
+            state = self.step(state, input_ids[:, t], W)
+            outs.append(self.logits_from_state(state))
+        logits = torch.stack(outs, 1)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1))
+        return FlyGPTOutput(loss=loss, logits=logits, state=state)
+
+    # ---- generation: carry the neuron state instead of a KV cache ------------------------------
+    @classmethod
+    def _supports_default_dynamic_cache(cls) -> bool:
+        return False  # stateful recurrent model: no KV cache, the neuron state is carried in `state`
+
+    def prepare_inputs_for_generation(self, input_ids, state=None, **kwargs):
+        if state is not None:
+            input_ids = input_ids[:, -1:]
+        return {"input_ids": input_ids, "state": state}
+
+    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, **kwargs):
+        model_kwargs["state"] = outputs.state
+        return model_kwargs

init/tokenizer.json

@@ -0,0 +1,90 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [],
+  "normalizer": null,
+  "pre_tokenizer": {
+    "type": "Split",
+    "pattern": {
+      "Regex": "[\\s\\S]"
+    },
+    "behavior": "Isolated",
+    "invert": false
+  },
+  "post_processor": null,
+  "decoder": {
+    "type": "Fuse"
+  },
+  "model": {
+    "type": "WordLevel",
+    "vocab": {
+      "\n": 0,
+      " ": 1,
+      "!": 2,
+      "$": 3,
+      "&": 4,
+      "'": 5,
+      ",": 6,
+      "-": 7,
+      ".": 8,
+      "3": 9,
+      ":": 10,
+      ";": 11,
+      "?": 12,
+      "A": 13,
+      "B": 14,
+      "C": 15,
+      "D": 16,
+      "E": 17,
+      "F": 18,
+      "G": 19,
+      "H": 20,
+      "I": 21,
+      "J": 22,
+      "K": 23,
+      "L": 24,
+      "M": 25,
+      "N": 26,
+      "O": 27,
+      "P": 28,
+      "Q": 29,
+      "R": 30,
+      "S": 31,
+      "T": 32,
+      "U": 33,
+      "V": 34,
+      "W": 35,
+      "X": 36,
+      "Y": 37,
+      "Z": 38,
+      "a": 39,
+      "b": 40,
+      "c": 41,
+      "d": 42,
+      "e": 43,
+      "f": 44,
+      "g": 45,
+      "h": 46,
+      "i": 47,
+      "j": 48,
+      "k": 49,
+      "l": 50,
+      "m": 51,
+      "n": 52,
+      "o": 53,
+      "p": 54,
+      "q": 55,
+      "r": 56,
+      "s": 57,
+      "t": 58,
+      "u": 59,
+      "v": 60,
+      "w": 61,
+      "x": 62,
+      "y": 63,
+      "z": 64
+    },
+    "unk_token": "<unk>"
+  }
+}

init/tokenizer_config.json

@@ -0,0 +1,6 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 1000000,
+  "clean_up_tokenization_spaces": false,
+  "add_prefix_space": false
+}