Albatross/rollout.py at main · RhineAI/Albatross · GitHub

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
import math, torch, types, copy, re
import numpy as np
from torch.nn import functional as F

args = types.SimpleNamespace()
args.vocab_size = 65536
args.head_size = 64
#
# model download: https://huggingface.co/BlinkDL/rwkv7-g1
#
# args.MODEL_NAME = "/mnt/e/RWKV-Runner/models/rwkv7-g1a-0.1b-20250728-ctx4096"
args.MODEL_NAME = "/mnt/e/RWKV-Runner/models/rwkv7-g1a-0.4b-20250905-ctx4096"
# args.MODEL_NAME = "/mnt/e/RWKV-Runner/models/rwkv7-g1-1.5b-20250429-ctx4096"
# args.MODEL_NAME = "/mnt/e/RWKV-Runner/models/rwkv7-g1-2.9b-20250519-ctx4096"
# args.MODEL_NAME = "/mnt/e/RWKV-Runner/models/rwkv7-g0a-7.2b-20250829-ctx4096"

prompt = 'User: Evaluate $(1+2i)6-3i$.\n\nAssistant: <think'
BATCH_SIZE = 320
GENERATION_LENGTH = 4000
BATCH_SIZE = 64
GENERATION_LENGTH = 1000
# BATCH_SIZE = 4
# GENERATION_LENGTH = 100

# we use simple sampling = greedy(logits + noise)
DECODE_NOISE = 1.0
DECODE_TEMP = 0.5
LOG_FILE = open("rollout.log", "w")

########################################################################################################

from reference.rwkv7 import RWKV_x070
from reference.utils import TRIE_TOKENIZER, sampler_simple_batch

# init model
print('loading...', args.MODEL_NAME)
model = RWKV_x070(args)
tokenizer = TRIE_TOKENIZER("reference/rwkv_vocab_v20230424.txt")

# init state
state = model.generate_zero_state(BATCH_SIZE)
out = model.forward_batch([tokenizer.encode(prompt) for _ in range(BATCH_SIZE)], state)

all_out = []

print(f'rollout {GENERATION_LENGTH} tokens...', end=' ')

for i in range(GENERATION_LENGTH):
    token = sampler_simple_batch(out, noise=DECODE_NOISE, temp=DECODE_TEMP).tolist()
    all_out.append(token)
    out = model.forward_batch(token, state)
    if i % 10 == 0:
        print(i, end=' ', flush=True)
print('\n' + '#'*120 + '\n')

all_out = np.transpose(np.array(all_out), axes=(1,0,2)).squeeze(-1)

for n in range(BATCH_SIZE):
    tokens = all_out[n]

    eod = np.flatnonzero(tokens == 0)
    if eod.size:
        tokens = tokens[:eod[0]] # get tokens before eod (token 0)

    out_str = tokenizer.decode(tokens, utf8_errors="ignore").strip()
    LOG_FILE.write(out_str + '\n' + '#'*120 + '\n')
    if eod.size:
        print(out_str.splitlines()[-1].strip())
    else:
        print(f'(unfinished within {GENERATION_LENGTH} tokens)')
    print('#'*120)

LOG_FILE.close()