llama

1
#include "arg.h"
2
#include "common.h"
3
#include "llama.h"
4

5
#include <vector>
6
#include <cstdio>
7

8
int main(int argc, char ** argv) {
9
    gpt_params params;
10

11
    params.prompt = "The quick brown fox";
12
    params.sparams.seed = 1234;
13

14
    if (!gpt_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
15
        return 1;
16
    }
17

18
    print_build_info();
19

20
    if (params.n_predict < 0) {
21
        params.n_predict = 16;
22
    }
23

24
    auto n_past = 0;
25

26
    std::string result0;
27
    std::string result1;
28
    std::string result2;
29

30
    // init
31
    llama_init_result llama_init = llama_init_from_gpt_params(params);
32

33
    llama_model * model = llama_init.model;
34
    llama_context * ctx = llama_init.context;
35

36
    if (model == nullptr || ctx == nullptr) {
37
        fprintf(stderr, "%s : failed to init\n", __func__);
38
        return 1;
39
    }
40

41
    auto sparams = llama_sampler_chain_default_params();
42

43
    llama_sampler * smpl = llama_sampler_chain_init(sparams);
44

45
    llama_sampler_chain_add(smpl, llama_sampler_init_softmax());
46
    llama_sampler_chain_add(smpl, llama_sampler_init_dist(params.sparams.seed));
47

48
    // tokenize prompt
49
    auto tokens = llama_tokenize(ctx, params.prompt, true);
50

51
    // evaluate prompt
52
    llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size(), n_past, 0));
53
    n_past += tokens.size();
54

55
    // save state (rng, logits, embedding and kv_cache) to file
56
    {
57
        std::vector<uint8_t> state_mem(llama_state_get_size(ctx));
58
        const size_t written = llama_state_get_data(ctx, state_mem.data(), state_mem.size());
59

60
        FILE *fp_write = fopen("dump_state.bin", "wb");
61
        fwrite(state_mem.data(), 1, written, fp_write);
62
        fclose(fp_write);
63

64
        fprintf(stderr, "%s : serialized state into %zd out of a maximum of %zd bytes\n", __func__, written, state_mem.size());
65
    }
66

67
    // save state (last tokens)
68
    const auto n_past_saved = n_past;
69

70
    // first run
71
    printf("\nfirst run: %s", params.prompt.c_str());
72

73
    for (auto i = 0; i < params.n_predict; i++) {
74
        auto next_token     = llama_sampler_sample(smpl, ctx, -1);
75
        auto next_token_str = llama_token_to_piece(ctx, next_token);
76

77
        printf("%s", next_token_str.c_str());
78
        result0 += next_token_str;
79

80
        if (llama_decode(ctx, llama_batch_get_one(&next_token, 1, n_past, 0))) {
81
            fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
82
            llama_free(ctx);
83
            llama_free_model(model);
84
            return 1;
85
        }
86
        n_past += 1;
87
    }
88

89
    printf("\n\n");
90

91
    // free old context
92
    llama_free(ctx);
93

94
    // make new context
95
    auto * ctx2 = llama_new_context_with_model(model, llama_context_params_from_gpt_params(params));
96

97
    llama_sampler * smpl2 = llama_sampler_chain_init(sparams);
98

99
    llama_sampler_chain_add(smpl2, llama_sampler_init_softmax());
100
    llama_sampler_chain_add(smpl2, llama_sampler_init_dist(params.sparams.seed));
101

102
    printf("\nsecond run: %s", params.prompt.c_str());
103

104
    // load state (rng, logits, embedding and kv_cache) from file
105
    {
106
        std::vector<uint8_t> state_mem;
107

108
        FILE * fp_read = fopen("dump_state.bin", "rb");
109
        fseek(fp_read, 0, SEEK_END);
110
        state_mem.resize(ftell(fp_read));
111
        fseek(fp_read, 0, SEEK_SET);
112
        const size_t read = fread(state_mem.data(), 1, state_mem.size(), fp_read);
113
        fclose(fp_read);
114

115
        if (read != llama_state_set_data(ctx2, state_mem.data(), state_mem.size())) {
116
            fprintf(stderr, "\n%s : failed to read state\n", __func__);
117
            llama_free(ctx2);
118
            llama_free_model(model);
119
            return 1;
120
        }
121

122
        fprintf(stderr, "%s : deserialized state from %zd out of a maximum of %zd bytes\n", __func__, read, state_mem.size());
123
    }
124

125
    // restore state (last tokens)
126
    n_past = n_past_saved;
127

128
    // second run
129
    for (auto i = 0; i < params.n_predict; i++) {
130
        auto next_token     = llama_sampler_sample(smpl2, ctx2, -1);
131
        auto next_token_str = llama_token_to_piece(ctx2, next_token);
132

133
        printf("%s", next_token_str.c_str());
134
        result1 += next_token_str;
135

136
        if (llama_decode(ctx2, llama_batch_get_one(&next_token, 1, n_past, 0))) {
137
            fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
138
            llama_free(ctx2);
139
            llama_free_model(model);
140
            return 1;
141
        }
142
        n_past += 1;
143
    }
144

145
    printf("\n\n");
146

147
    llama_free(ctx2);
148

149
    if (result0 != result1) {
150
        fprintf(stderr, "\n%s : error : the 2 generations are different\n", __func__);
151
        return 1;
152
    }
153

154
    // make new context
155
    auto * ctx3 = llama_new_context_with_model(model, llama_context_params_from_gpt_params(params));
156

157
    llama_sampler * smpl3 = llama_sampler_chain_init(sparams);
158

159
    llama_sampler_chain_add(smpl3, llama_sampler_init_softmax());
160
    llama_sampler_chain_add(smpl3, llama_sampler_init_dist(params.sparams.seed));
161

162
    printf("\nsingle seq run: %s", params.prompt.c_str());
163

164
    // load state (rng, logits, embedding and kv_cache) from file
165
    {
166
        std::vector<uint8_t> state_mem;
167

168
        FILE * fp_read = fopen("dump_state.bin", "rb");
169
        fseek(fp_read, 0, SEEK_END);
170
        state_mem.resize(ftell(fp_read));
171
        fseek(fp_read, 0, SEEK_SET);
172
        const size_t read = fread(state_mem.data(), 1, state_mem.size(), fp_read);
173
        fclose(fp_read);
174

175
        if (read != llama_state_set_data(ctx3, state_mem.data(), state_mem.size())) {
176
            fprintf(stderr, "\n%s : failed to read state\n", __func__);
177
            llama_free(ctx3);
178
            llama_free_model(model);
179
            return 1;
180
        }
181

182
        fprintf(stderr, "%s : deserialized state from %zd out of a maximum of %zd bytes\n", __func__, read, state_mem.size());
183
    }
184

185
    // restore state (last tokens)
186
    n_past = n_past_saved;
187

188
    // save seq 0 and load into seq 1
189
    {
190
        // save kv of seq 0
191
        std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx3, 0));
192
        const size_t ncopy = llama_state_seq_get_data(ctx3, seq_store.data(), seq_store.size(), 0);
193
        if (ncopy != seq_store.size()) {
194
            fprintf(stderr, "\n%s : seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
195
            llama_free(ctx3);
196
            llama_free_model(model);
197
            return 1;
198
        }
199
        fprintf(stderr, "%s : seq 0 copied, %zd bytes\n", __func__, ncopy);
200

201
        // erase whole kv
202
        llama_kv_cache_clear(ctx3);
203
        fprintf(stderr, "%s : kv cache cleared\n", __func__);
204

205
        // restore kv into seq 1
206
        const size_t nset = llama_state_seq_set_data(ctx3, seq_store.data(), seq_store.size(), 1);
207
        if (nset != seq_store.size()) {
208
            fprintf(stderr, "\n%s : seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
209
            llama_free(ctx3);
210
            llama_free_model(model);
211
            return 1;
212
        }
213
        fprintf(stderr, "%s : seq 1 restored, %zd bytes\n", __func__, nset);
214
    }
215

216
    // third run with seq 1 instead of 0
217
    for (auto i = 0; i < params.n_predict; i++) {
218
        auto next_token     = llama_sampler_sample(smpl3, ctx3, -1);
219
        auto next_token_str = llama_token_to_piece(ctx3, next_token);
220

221
        printf("%s", next_token_str.c_str());
222
        result2 += next_token_str;
223

224
        if (llama_decode(ctx3, llama_batch_get_one(&next_token, 1, n_past, 1))) {
225
            fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
226
            llama_free(ctx3);
227
            llama_free_model(model);
228
            return 1;
229
        }
230
        n_past += 1;
231
    }
232

233
    printf("\n");
234

235
    llama_sampler_free(smpl);
236
    llama_sampler_free(smpl2);
237
    llama_sampler_free(smpl3);
238

239
    llama_free(ctx3);
240
    llama_free_model(model);
241

242
    if (result0 != result2) {
243
        fprintf(stderr, "\n%s : error : the seq restore generation is different\n", __func__);
244
        return 1;
245
    }
246

247
    fprintf(stderr, "\n%s : success\n", __func__);
248

249
    return 0;
250
}
251