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