README.md

---
language: en
license: apache-2.0
library_name: transformers
pipeline_tag: text-generation
tags:
  - interactive-fiction
  - text-adventure
  - storytelling
  - custom_code
---

![Bard Small](banner.png)

<p align="center">
  <a href="https://slop.szymon.space">Play in the browser</a> ·
  <a href="https://huggingface.co/spaces/lowerkey/stochastic-lands-of-prophecy">Space</a> ·
  <a href="#using-it">Use it from Python</a>
</p>

Bard Small is a 137M parameter model that runs text adventures. You type a first line, it writes what happens next,
tells the game what changed (you lost some money, picked up a key) and gives you a few choices. Pick one or type
something else and it keeps going.

I trained it from scratch for a game I made, [Stochastic Lands of Prophecy](https://slop.szymon.space). The game runs
the whole model in your browser, so there's no server and nothing you type leaves your machine. You can try it there
or in the [Space](https://huggingface.co/spaces/lowerkey/stochastic-lands-of-prophecy). This repo has the weights plus
the code that goes around them (tokenizer, number handling, the game rules) so you can also use it from Python.

## A sample

The opening was `You arrive at the gates of the old castle.` and I picked the first choice. I ran it a few times and
took one of the better ones.

```
You arrive at the gates of the old castle. “You’ve come a long way, little Prince. The guards say that you are only
here for a few days. If you wish for it, I can take you back home and live somewhere else. Just tell me where your
home is and the guards will let you pass.”

This request really brings a smile to your face. The king asks you to accompany him to his manor and there you will
learn more.

The gate opens and you follow him inside. He takes a deep breath before leading you into a room. You see a man sitting
on a chair.

“I have a question for you son,” he says with a smile.

  1. “What do you want to ask me?”
  2. “What will you be doing now?”
  3. “What is your name?”

> “What do you want to ask me?”

“Well, Prince, I have been summoned to King Arthur’s castle and I have come to ask for your help. Do you know where
your home is located?”

You think about this for a moment. The king is waiting for you to speak but he already knows.
```

It loses the thread pretty quickly. More on that at the bottom.

## Using it

```python
# pip install torch transformers tokenmonster
from transformers import AutoModelForCausalLM

bard = AutoModelForCausalLM.from_pretrained("lowerkey/bard-small", trust_remote_code=True)

tale = bard.begin("You wake up in a dungeon.")
print(tale.text)       # what the model wrote
print(tale.choices)    # what you can do next

tale = bard.act(tale, tale.choices[0])                  # take a choice
tale = bard.act(tale, "I ask the guard about the key")  # or type your own
```

`bard.play()` runs the same loop in the terminal if you only want to poke at it.

`begin` takes a few extra arguments. `length` is "short", "medium" or "long" and controls how much it writes per turn.
`state` is what the player carries and defaults to `{"money": 100}`. `seed` makes a run repeatable. `on_text` is a
function that receives the text in pieces while it's being written, which is handy for streaming.

A `Tale` is a small dataclass and the whole story is in `tale.history`, a plain string. To save one use
`tale.to_dict()`, and to load it again `Tale.from_dict(d)`.

On a GPU load it in bf16 with `from_pretrained(..., dtype=torch.bfloat16).to("cuda")` (the argument is called
`torch_dtype` before transformers 4.56). A turn takes about a second on a GPU and a few seconds on a CPU. I tried it
with transformers 4.46 and 5.19.

You need `trust_remote_code=True` because the architecture isn't part of transformers. The code it runs is the `.py`
files in this repo.

## Why there's no generate()

Numbers don't go through the tokenizer. Every run of digits in a text turns into one `<|num|>` token, and the actual
value travels in a second stream next to the token ids. A small network turns that value into a vector that gets added
to the token's embedding, and when the model wants to write a number a second small network predicts its digits. The
model's input is therefore two tensors, `input_ids` and `number_values`, and the usual `generate()` only knows about
the first. `bard.begin` and `bard.act` take care of this. If you want raw logits, `bard(input_ids, number_values)`
gives you those.

## The format

The model reads and writes tales in a small text format. A tale is a list of blocks, and every block is

```
<|start|>ROLE<|content|>text<|end|>
```

where ROLE is `state`, `story` or `player`. The role headers and all the other markers are single tokens.

A `state` block says what the player has. Each variable is its name, a newline and its value, and variables are
separated by `<|sep|>`. The game writes this block, not the model, and puts a new one in after every turn.

A `player` block is the player's move: the text of the choice they picked or whatever they typed. Also written by the
game.

A `story` block is the only one the model writes. First comes the narration. After that, if something changed,
`<|vars|>` and then the changes separated by `<|sep|>`, each one a variable name, then `<|add|>` or `<|set|>`, then
a number. Last comes `<|choices|>` and the choices, again separated by `<|sep|>`. When the model writes a story block
with no choices, the tale is over.

Put together it looks like this (the passage is shortened):

```
<|start|>state<|content|>money
100<|end|>
<|start|>story<|content|>The guard looks at your coin purse and smiles.<|vars|>money<|add|>-20<|choices|>Pay him<|sep|>Run<|end|>
<|start|>state<|content|>money
80<|end|>
<|start|>player<|content|>Pay him<|end|>
<|start|>story<|content|>He steps aside and the door swings open. ...
```

What the model puts in `<|vars|>` is a suggestion. A model this small will cheerfully write `money -400` when you
have 100, or take your money after a fight where nobody had any. `bard_rules.py` looks at each change before it goes
through. A change to money has to be backed up by the passage actually talking about money, an amount can't drop below
zero or grow tenfold in one go, and so on. `bard_steering.py` keeps passages from running on forever and makes sure
there are at least two different choices.

## The model

It has 36 layers, a width of 512 and 8 attention heads, for 136.7M parameters. Attention and the feed forward layer run
in parallel in each block. Most layers look at a sliding window of 256 tokens with RoPE, and every fourth layer looks
at everything with no position encoding. There's QK-norm and a gate on the attention output, and the logits are
softcapped. The context is 4096 tokens, and older text falls off the back as a tale gets long.

The vocabulary has 16,128 ids: tokenmonster's `fiction-16000-strict-v1` plus the special tokens above. The weights are
in bf16 (`model.safetensors`, 273 MB).

## Training

Three stages, all from scratch. First plain English, about 11B tokens of FineWeb-Edu. Then fiction: Royal Road
stories, Project Gutenberg fiction, WritingPrompts and BookCorpus, with some FineWeb mixed back in so it doesn't forget
how to write normal sentences. Last, choose-your-own-adventure stories converted into the block format above, which is
where it learns to write narration, changes and choices in that shape.

## What it's bad at

Plenty. It's small. Names and details change between turns, the plot wanders off, and now and then a turn makes no
sense at all. It reads more like a dream than a story. If it gets stuck, pick a different choice or type your own
action, that usually shakes it loose.

It was trained on web text and fiction, some of it for adults, and I did no safety tuning. It can write violent or
sexual things, and it can be biased. Don't use it for anything important.

## License and data

The weights and the code here are Apache 2.0. The training data is publicly available text. Some of it, Royal Road and
BookCorpus in particular, belongs to its authors and has its own terms, so if you want to build something on these
weights, check that it's fine for what you have in mind. The tokenizer vocabulary comes from tokenmonster's
`fiction-16000-strict-v1` (MIT).