Neuraxis-stories

Custom ~270M-parameter Archon transformer trained for story generation. Tokenization uses GPT-2 BPE via tiktoken.

Files

File	Description
Neuraxis.pt	Model weights (state_dict)
config.json	Architecture hyperparameters
loss_history.json	Train / validation loss curves

Quick start

import torch
import tiktoken
from huggingface_hub import hf_hub_download

from archon.model import ArchonModel
from archon.config import load_config

repo_id = "viratarun/Neuraxis-stories"
device = "cuda" if torch.cuda.is_available() else "cpu"

config_path = hf_hub_download(repo_id, "config.json")
weights_path = hf_hub_download(repo_id, "Neuraxis.pt")

model = ArchonModel(load_config(config_path))
model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
model.to(device).eval()

enc = tiktoken.get_encoding("gpt2")
prompt = "A little girl went to the woods"
context = torch.tensor(enc.encode_ordinary(prompt)).unsqueeze(0).to(device)

with torch.no_grad():
    out = model.generate(context, max_new_tokens=200, temperature=0.8, top_k=50)

print(enc.decode(out.squeeze().tolist()))

Architecture

• Embedding dim: 640
• Layers: 18 (sliding + full attention)
• Heads: 4 (GQA, 1 KV group)
• Context length: 1024
• Vocab: 50257 (GPT-2)