Tallow / tallow.py

add tallow.py

69092c3 verified 2 days ago

6.97 kB

"""
Talk to Tallow, a 15.7M-parameter crow.
    pip install torch sentencepiece safetensors
    python tallow.py
Type to him. Empty line = quit. /new = start over as a stranger.
He sees the whole conversation every turn, so he remembers how you treated him (within ~1,900 tokens).
Everything needed is in this one file: the model, and the two text tricks his tokenizer expects
(a <nl> token for line breaks, and a space after every OPENING "*" or "[").
"""
import json
import math
import os
from dataclasses import dataclass

import sentencepiece as spm
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file

HERE = os.path.dirname(os.path.abspath(__file__))
DEV = "cuda" if torch.cuda.is_available() else "cpu"


# ---------------------------------------------------------------- the model (Llama-style, written out by hand)
@dataclass
class Config:
    vocab: int = 4096
    ctx: int = 2048
    d: int = 384
    layers: int = 8
    heads: int = 6
    hidden: int = 1024


class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        xf = x.float()
        return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight).type_as(x)


def rope_tables(head_dim, ctx, base=10000.0):
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    ang = torch.outer(torch.arange(ctx).float(), inv)
    ang = torch.cat([ang, ang], dim=-1)
    return ang.cos(), ang.sin()


def rotate(x, cos, sin):
    x1, x2 = x.chunk(2, dim=-1)
    return x * cos + torch.cat([-x2, x1], dim=-1) * sin


class Attention(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.heads, self.hd = c.heads, c.d // c.heads
        self.qkv = nn.Linear(c.d, 3 * c.d, bias=False)
        self.out = nn.Linear(c.d, c.d, bias=False)

    def forward(self, x, cos, sin):
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=2)
        q, k, v = (t.view(B, T, self.heads, self.hd).transpose(1, 2) for t in (q, k, v))
        q, k = rotate(q, cos, sin), rotate(k, cos, sin)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(a.transpose(1, 2).reshape(B, T, D))


class SwiGLU(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.gate = nn.Linear(c.d, c.hidden, bias=False)
        self.up = nn.Linear(c.d, c.hidden, bias=False)
        self.down = nn.Linear(c.hidden, c.d, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.norm1, self.attn = RMSNorm(c.d), Attention(c)
        self.norm2, self.mlp = RMSNorm(c.d), SwiGLU(c)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        return x + self.mlp(self.norm2(x))


class Tallow(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.c = c
        self.embed = nn.Embedding(c.vocab, c.d)
        self.blocks = nn.ModuleList(Block(c) for _ in range(c.layers))
        self.norm = RMSNorm(c.d)
        cos, sin = rope_tables(c.d // c.heads, c.ctx)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def forward(self, idx):
        T = idx.shape[1]
        x = self.embed(idx)
        cos, sin = self.cos[:T].to(x.dtype), self.sin[:T].to(x.dtype)
        for b in self.blocks:
            x = b(x, cos, sin)
        return self.norm(x) @ self.embed.weight.T      # tied embeddings


# ---------------------------------------------------------------- the tokenizer and its two text tricks
NL = "<nl>"


def to_sp(text):
    """Line breaks -> <nl>; a space after every OPENING '*' (1st, 3rd... on a line) and every '['."""
    lines = []
    for line in text.split("\n"):
        out, stars = [], 0
        for i, ch in enumerate(line):
            out.append(ch)
            stars += ch == "*"
            opener = ch == "[" or (ch == "*" and stars % 2 == 1)
            if opener and i + 1 < len(line) and not line[i + 1].isspace():
                out.append(" ")
        lines.append("".join(out))
    return (NL + " ").join(lines)


def from_sp(text):
    text = text.replace(NL + " ", "\n").replace(NL, "\n")
    lines = []
    for line in text.split("\n"):
        out, stars, skip = [], 0, False
        for i, ch in enumerate(line):
            if skip:
                skip = False
                continue
            out.append(ch)
            stars += ch == "*"
            opener = ch == "[" or (ch == "*" and stars % 2 == 1)
            if opener and line[i + 1:i + 2] == " " and i + 2 < len(line) and not line[i + 2].isspace():
                skip = True
        lines.append("".join(out))
    return "\n".join(lines)


def render(sp, history):
    """you: ...<nl>tallow: ...<nl>  (the exact format he was trained on)"""
    nl = sp.piece_to_id(NL)
    ids = []
    for turn in history:
        if turn["role"] == "user":
            ids += sp.encode(to_sp(f"you: {turn['text']}")) + [nl]
        else:
            ids += sp.encode(to_sp("tallow:")) + sp.encode(to_sp(turn["text"])) + [nl]
    return ids


@torch.no_grad()
def reply(model, sp, history, temperature=0.8, top_k=50, max_new=80):
    while len(history) > 1 and len(render(sp, history)) > 1900:   # forget the oldest lines if too long
        history = history[2:]
    ids = render(sp, history) + sp.encode(to_sp("tallow:"))
    start, stop = len(ids), {sp.piece_to_id(NL), sp.eos_id()}
    for _ in range(max_new):
        logits = model(torch.tensor([ids[-model.c.ctx:]], device=DEV))[0, -1].float() / temperature
        logits[logits < torch.topk(logits, top_k).values[-1]] = -float("inf")
        nxt = int(torch.multinomial(F.softmax(logits, -1), 1))
        if nxt in stop:
            break
        ids.append(nxt)
    return from_sp(sp.decode(ids[start:])).strip()


def load():
    with open(os.path.join(HERE, "config.json"), encoding="utf-8") as f:
        model = Tallow(Config(**json.load(f)))
    model.load_state_dict(load_file(os.path.join(HERE, "model.safetensors")))
    model.to(DEV).eval()
    return model, spm.SentencePieceProcessor(model_file=os.path.join(HERE, "tallow_tok.model"))


if __name__ == "__main__":
    model, sp = load()
    history = []
    print("Tallow is on the power line. (empty line = quit, /new = start over)\n")
    while True:
        you = input("you: ").strip()
        if not you:
            break
        if you == "/new":
            history = []
            print("(he doesn't know you anymore)\n")
            continue
        history.append({"role": "user", "text": you})
        line = reply(model, sp, history)
        history.append({"role": "tallow", "text": line})
        print(f"tallow: {line}\n")

	"""
	Talk to Tallow, a 15.7M-parameter crow.

	pip install torch sentencepiece safetensors
	python tallow.py

	Type to him. Empty line = quit. /new = start over as a stranger.
	He sees the whole conversation every turn, so he remembers how you treated him (within ~1,900 tokens).

	Everything needed is in this one file: the model, and the two text tricks his tokenizer expects
	(a <nl> token for line breaks, and a space after every OPENING "*" or "[").
	"""
	import json
	import math
	import os
	from dataclasses import dataclass

	import sentencepiece as spm
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from safetensors.torch import load_file

	HERE = os.path.dirname(os.path.abspath(__file__))
	DEV = "cuda" if torch.cuda.is_available() else "cpu"


	# ---------------------------------------------------------------- the model (Llama-style, written out by hand)
	@dataclass
	class Config:
	vocab: int = 4096
	ctx: int = 2048
	d: int = 384
	layers: int = 8
	heads: int = 6
	hidden: int = 1024


	class RMSNorm(nn.Module):
	def __init__(self, d, eps=1e-5):
	super().__init__()
	self.weight = nn.Parameter(torch.ones(d))
	self.eps = eps

	def forward(self, x):
	xf = x.float()
	return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight).type_as(x)


	def rope_tables(head_dim, ctx, base=10000.0):
	inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
	ang = torch.outer(torch.arange(ctx).float(), inv)
	ang = torch.cat([ang, ang], dim=-1)
	return ang.cos(), ang.sin()


	def rotate(x, cos, sin):
	x1, x2 = x.chunk(2, dim=-1)
	return x * cos + torch.cat([-x2, x1], dim=-1) * sin


	class Attention(nn.Module):
	def __init__(self, c):
	super().__init__()
	self.heads, self.hd = c.heads, c.d // c.heads
	self.qkv = nn.Linear(c.d, 3 * c.d, bias=False)
	self.out = nn.Linear(c.d, c.d, bias=False)

	def forward(self, x, cos, sin):
	B, T, D = x.shape
	q, k, v = self.qkv(x).split(D, dim=2)
	q, k, v = (t.view(B, T, self.heads, self.hd).transpose(1, 2) for t in (q, k, v))
	q, k = rotate(q, cos, sin), rotate(k, cos, sin)
	a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
	return self.out(a.transpose(1, 2).reshape(B, T, D))


	class SwiGLU(nn.Module):
	def __init__(self, c):
	super().__init__()
	self.gate = nn.Linear(c.d, c.hidden, bias=False)
	self.up = nn.Linear(c.d, c.hidden, bias=False)
	self.down = nn.Linear(c.hidden, c.d, bias=False)

	def forward(self, x):
	return self.down(F.silu(self.gate(x)) * self.up(x))


	class Block(nn.Module):
	def __init__(self, c):
	super().__init__()
	self.norm1, self.attn = RMSNorm(c.d), Attention(c)
	self.norm2, self.mlp = RMSNorm(c.d), SwiGLU(c)

	def forward(self, x, cos, sin):
	x = x + self.attn(self.norm1(x), cos, sin)
	return x + self.mlp(self.norm2(x))


	class Tallow(nn.Module):
	def __init__(self, c):
	super().__init__()
	self.c = c
	self.embed = nn.Embedding(c.vocab, c.d)
	self.blocks = nn.ModuleList(Block(c) for _ in range(c.layers))
	self.norm = RMSNorm(c.d)
	cos, sin = rope_tables(c.d // c.heads, c.ctx)
	self.register_buffer("cos", cos, persistent=False)
	self.register_buffer("sin", sin, persistent=False)

	def forward(self, idx):
	T = idx.shape[1]
	x = self.embed(idx)
	cos, sin = self.cos[:T].to(x.dtype), self.sin[:T].to(x.dtype)
	for b in self.blocks:
	x = b(x, cos, sin)
	return self.norm(x) @ self.embed.weight.T # tied embeddings


	# ---------------------------------------------------------------- the tokenizer and its two text tricks
	NL = "<nl>"


	def to_sp(text):
	"""Line breaks -> <nl>; a space after every OPENING '*' (1st, 3rd... on a line) and every '['."""
	lines = []
	for line in text.split("\n"):
	out, stars = [], 0
	for i, ch in enumerate(line):
	out.append(ch)
	stars += ch == "*"
	opener = ch == "[" or (ch == "*" and stars % 2 == 1)
	if opener and i + 1 < len(line) and not line[i + 1].isspace():
	out.append(" ")
	lines.append("".join(out))
	return (NL + " ").join(lines)


	def from_sp(text):
	text = text.replace(NL + " ", "\n").replace(NL, "\n")
	lines = []
	for line in text.split("\n"):
	out, stars, skip = [], 0, False
	for i, ch in enumerate(line):
	if skip:
	skip = False
	continue
	out.append(ch)
	stars += ch == "*"
	opener = ch == "[" or (ch == "*" and stars % 2 == 1)
	if opener and line[i + 1:i + 2] == " " and i + 2 < len(line) and not line[i + 2].isspace():
	skip = True
	lines.append("".join(out))
	return "\n".join(lines)


	def render(sp, history):
	"""you: ...<nl>tallow: ...<nl> (the exact format he was trained on)"""
	nl = sp.piece_to_id(NL)
	ids = []
	for turn in history:
	if turn["role"] == "user":
	ids += sp.encode(to_sp(f"you: {turn['text']}")) + [nl]
	else:
	ids += sp.encode(to_sp("tallow:")) + sp.encode(to_sp(turn["text"])) + [nl]
	return ids


	@torch.no_grad()
	def reply(model, sp, history, temperature=0.8, top_k=50, max_new=80):
	while len(history) > 1 and len(render(sp, history)) > 1900: # forget the oldest lines if too long
	history = history[2:]
	ids = render(sp, history) + sp.encode(to_sp("tallow:"))
	start, stop = len(ids), {sp.piece_to_id(NL), sp.eos_id()}
	for _ in range(max_new):
	logits = model(torch.tensor([ids[-model.c.ctx:]], device=DEV))[0, -1].float() / temperature
	logits[logits < torch.topk(logits, top_k).values[-1]] = -float("inf")
	nxt = int(torch.multinomial(F.softmax(logits, -1), 1))
	if nxt in stop:
	break
	ids.append(nxt)
	return from_sp(sp.decode(ids[start:])).strip()


	def load():
	with open(os.path.join(HERE, "config.json"), encoding="utf-8") as f:
	model = Tallow(Config(**json.load(f)))
	model.load_state_dict(load_file(os.path.join(HERE, "model.safetensors")))
	model.to(DEV).eval()
	return model, spm.SentencePieceProcessor(model_file=os.path.join(HERE, "tallow_tok.model"))


	if __name__ == "__main__":
	model, sp = load()
	history = []
	print("Tallow is on the power line. (empty line = quit, /new = start over)\n")
	while True:
	you = input("you: ").strip()
	if not you:
	break
	if you == "/new":
	history = []
	print("(he doesn't know you anymore)\n")
	continue
	history.append({"role": "user", "text": you})
	line = reply(model, sp, history)
	history.append({"role": "tallow", "text": line})
	print(f"tallow: {line}\n")