File size: 5,013 Bytes
dde517d 8414bb4 dde517d 8414bb4 dde517d 4f44f20 dde517d 8414bb4 dde517d 8414bb4 dde517d 8414bb4 dde517d | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """CBT-18.95M: character-level Transformer with Decoupled 2D-RoPE.
Quick start:
from cbt import load, add, chat
model, tok = load() # local folder or the Hugging Face repo
add(model, tok, 123456789, 987654321) # {'answer': '1111111110', 'correct': True, ...}
chat(model, tok, "What is the weather in Tokyo?")
twin, _ = load(variant="twin_1d") # matched 1D-RoPE control (same architecture, data, steps)
"""
import json
import os
from dataclasses import fields
from typing import Any, Callable, Dict, Optional, Tuple, Union
import torch
from .model import ScaledGPT as CBT, ScaledGPTConfig as CBTConfig
from .tokenizer import ScaledTokenizer as CBTTokenizer
from .tools import execute_tool
__all__ = ["CBT", "CBTConfig", "CBTTokenizer", "load", "add", "chat", "execute_tool"]
REPO_ID = "Neithanvarvas/CBT"
VARIANTS: Dict[str, str] = {"cbt_2d": "", "twin_1d": "twin_1d/"}
_PKG_DIR = os.path.dirname(os.path.abspath(__file__))
def _resolve(source: Optional[str], filename: str) -> str:
"""Finds `filename` in a local folder (default: the repo root next to this package) or downloads it."""
folder = source or os.path.dirname(_PKG_DIR)
local = os.path.join(folder, filename)
# A clone made without git-lfs leaves a ~130-byte pointer file instead of the weights.
if os.path.isfile(local) and not (filename.endswith(".safetensors") and os.path.getsize(local) < 1024):
return local
from huggingface_hub import hf_hub_download
return hf_hub_download(repo_id=source if source and not os.path.isdir(source) else REPO_ID, filename=filename)
def load(
source: Optional[str] = None,
device: Optional[Union[str, torch.device]] = None,
variant: str = "cbt_2d",
) -> Tuple[CBT, CBTTokenizer]:
"""Loads model and tokenizer from a local folder or a Hugging Face repo id; returns (model, tokenizer) in eval mode.
variant: "cbt_2d" (the released model) or "twin_1d" (its 1D-RoPE control).
"""
from safetensors.torch import load_file
if variant not in VARIANTS:
raise ValueError(f"variant must be one of {sorted(VARIANTS)}")
prefix = VARIANTS[variant]
with open(_resolve(source, prefix + "config.json"), encoding="utf-8") as f:
raw: Dict[str, Any] = json.load(f)
config = CBTConfig(**{k.name: raw[k.name] for k in fields(CBTConfig) if k.name in raw})
device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
tok = CBTTokenizer()
model = CBT(config, tokenizer=tok)
model.load_state_dict(load_file(_resolve(source, prefix + "model.safetensors")))
return model.to(device).eval(), tok
@torch.no_grad()
def _generate(model: CBT, ids: torch.Tensor, max_new_tokens: int, stop: list) -> torch.Tensor:
return model.generate(ids, max_new_tokens=max_new_tokens, temperature=0.0, stop_tokens=stop)
def _encode(model: CBT, tok: CBTTokenizer, text: str, bos: bool = True) -> torch.Tensor:
device = next(model.parameters()).device
return torch.tensor([tok.encode(text, add_bos=bos, add_eos=False)], dtype=torch.long, device=device)
def add(model: CBT, tok: CBTTokenizer, a: Union[int, str], b: Union[int, str]) -> Dict[str, Any]:
"""Greedy-decodes a + b with the column scratchpad. Trained on 1-25 digit operands."""
a, b = str(a).strip(), str(b).strip()
if not (a.isdigit() and b.isdigit()):
raise ValueError("operands must be non-negative integers")
budget = 12 * (max(len(a), len(b)) + 2) + 8
out = _generate(model, _encode(model, tok, f"{a}+{b}=<THINK>"), budget, [tok.eos_id])
raw = tok.decode(out[0].tolist())
think, _, tail = raw.partition("<THINK>")[2].partition("</THINK>")
answer = tail.lstrip("=").replace("<EOS>", "").strip()
return {"answer": answer, "scratchpad": think, "correct": answer == str(int(a) + int(b)), "raw": raw}
def chat(
model: CBT,
tok: CBTTokenizer,
message: str,
tool_executor: Optional[Callable[[str], str]] = execute_tool,
max_new_tokens: int = 256,
) -> Dict[str, Any]:
"""Runs one user turn: <THINK> reasoning, optional <CALL> executed by `tool_executor`, then <ANSWER>."""
ids = _encode(model, tok, f"<USER>{message}</USER><THINK>")
out = _generate(model, ids, max_new_tokens, [tok.call_e_id, tok.answer_e_id, tok.eos_id])
raw = tok.decode(out[0].tolist())
if "</CALL>" in raw and tool_executor is not None:
call = raw.split("<CALL>")[-1].split("</CALL>")[0]
cont = torch.cat([out, _encode(model, tok, f"<RESULT>{tool_executor(call)}</RESULT>", bos=False)], dim=1)
out = _generate(model, cont, max_new_tokens, [tok.answer_e_id, tok.eos_id])
raw = tok.decode(out[0].tolist())
def span(tag: str) -> Optional[str]:
s, e = f"<{tag}>", f"</{tag}>"
return raw.split(s, 1)[1].split(e, 1)[0].strip() if s in raw and e in raw else None
return {"think": span("THINK"), "call": span("CALL"), "result": span("RESULT"), "answer": span("ANSWER"), "raw": raw}
|