| --- |
| license: apache-2.0 |
| tags: |
| - text-generation |
| - pytorch |
| - stories |
| library_name: pytorch |
| pipeline_tag: text-generation |
| --- |
| |
| # 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 |
|
|
| ```python |
| 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) |
|
|
