---
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)