# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "chipformer @ https://github.com/JGalego/olmputer/archive/98639dff4c84a18542e3ca12a286a4df563b3a27.tar.gz#subdirectory=experimental/riscv",  # pylint: disable=line-too-long
#     "huggingface_hub",
#     "torch==2.14.1",
#     "transformers==5.19.0",
# ]
# ///
"""Run fibonacci(25) on olmputer-tiny, the RISC-V computer in its weights.

The program's code and registers are written into the model's Gated DeltaNet states, and greedy
generate then runs the CPU: each token it generates is the next token of the execution trace. The
program ends with mv a0, t0 (the result), li a7, 93 and the exit ecall, so the result is the
second-to-last value record the model wrote.
"""

import torch
from chipformer.q35.ir import OLMO
from chipformer.rv import core, firmware, ref, spec
from chipformer.rv.programs import program
from transformers import OlmoHybridForCausalLM

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = OlmoHybridForCausalLM.from_pretrained("jgalego/olmputer-tiny", dtype=torch.bfloat16)
MODEL = MODEL.to(DEVICE).eval()
CIRCUIT = core.build(OLMO)  # the CPU as a circuit: it writes the program into the model's cache

TRACE = firmware.run(MODEL, CIRCUIT, ref.Machine.from_program(program("fib")))
VALUES = [spec.decode_value(TRACE[:i + 1]) for i, t in enumerate(TRACE)
          if spec.family(t)[0] in ("F0", "F1", "F2")]
print("fibonacci(25) =", VALUES[-2])
print(f"{len(TRACE)} trace tokens on {DEVICE}, identical to the reference machine:",
      TRACE == ref.Machine.from_program(program("fib")).run()[1:])