lancer/main_test.py at master · saturn-acc/lancer · 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
from pathlib import Path
from tokenizer import encode, get_input_embeds, decode
from attention import GroupedQueryAttention
from MultiLayerPerceptron import MLP
import mlx.core as mx
import mlx.nn as nn

text = input("what do you have to say:")
token_ids = encode(text)
embeds = get_input_embeds(token_ids)


class RMSnorm(nn.Module):
    def __init__(self,hidden_dim, load_gamma):
        super().__init__()
        self.norm = nn.RMSNorm(dims=hidden_dim, eps=1e-5)

        self.norm.weight = load_gamma

    def __call__(self, embeds):
         rmsn_embeds = self.norm(embeds)
         return rmsn_embeds


WEIGHTS = str(Path(__file__).parent / "llama_intruct3.2/snapshots/9213176726f574b556790deb65791e0c5aa438b6/model.safetensors")
weights = mx.load(WEIGHTS)

NUM_LAYERS = 16


next_id = None
eot_id = 128009

while next_id != eot_id:
    hidden = embeds
    for i in range(NUM_LAYERS):
        gamma = weights[f"model.layers.{i}.input_layernorm.weight"]
        rms = RMSnorm(hidden_dim=2048, load_gamma=gamma)
        rmsn_embeds = rms(hidden)

        Wq = weights[f"model.layers.{i}.self_attn.q_proj.weight"].T
        Wk = weights[f"model.layers.{i}.self_attn.k_proj.weight"].T
        Wv = weights[f"model.layers.{i}.self_attn.v_proj.weight"].T
        Wo = weights[f"model.layers.{i}.self_attn.o_proj.weight"].T

        attention = GroupedQueryAttention(Wk=Wk, Wq=Wq, Wv=Wv, Wo=Wo)
        att_out = attention(rmsn_embeds)
        hidden = hidden + att_out

        post_gamma = weights[f"model.layers.{i}.post_attention_layernorm.weight"]
        rms = RMSnorm(hidden_dim=2048, load_gamma=post_gamma)
        mlp_input = rms(hidden)

        W_gate = weights[f"model.layers.{i}.mlp.gate_proj.weight"].T
        W_up = weights[f"model.layers.{i}.mlp.up_proj.weight"].T
        W_down = weights[f"model.layers.{i}.mlp.down_proj.weight"].T

        mlp = MLP(W_gate=W_gate, W_up=W_up, W_down=W_down)
        mlp_output = mlp(mlp_input)
        hidden = hidden + mlp_output


    final_gamma = weights["model.norm.weight"]
    rms = RMSnorm(hidden_dim=2048, load_gamma=final_gamma)
    hidden = rms(hidden)

    vocab = weights["model.embed_tokens.weight"]
    logits = hidden @ vocab.T
    next_id = int(mx.argmax(logits[-1]))
    print(next_id, repr(decode([next_id])))
    new_row = vocab[next_id][None, :]
    embeds = mx.concatenate([embeds, new_row], axis=0)