File size: 1,481 Bytes
32f50e7 | 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 | ---
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)
|