char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aaron Teoargs: refactor mlock/mmap/directio into load-mode (#20834)e6dd0e29a

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.7 KiB266 linesraw
1#include "arg.h"
2#include "chat.h"
3#include "common.h"
4#include "diffusion.h"
5#include "llama.h"
6#include "log.h"
7
8#include <limits.h>
9
10#include <clocale>
11#include <cstring>
12#include <string>
13#include <vector>
14
15struct callback_data {
16    diffusion_params *  diff_params;
17    const llama_vocab * vocab;
18    int32_t             n_input;
19};
20
21static bool diffusion_step_callback(int32_t             step,
22                                    int32_t             total_steps,
23                                    const llama_token * tokens,
24                                    int32_t             n_tokens,
25                                    void *              user_data) {
26    (void) user_data;
27
28    callback_data * data = static_cast<callback_data *>(user_data);
29
30    auto print_progress_bar = [](int32_t step, int32_t total_steps) {
31        int progress_percent = (step * 100) / total_steps;
32        int progress_bars    = (step * 50) / total_steps;
33        LOG_INF("\rdiffusion step: %d/%d [%s%s] %d%%",
34                step,
35                total_steps,
36                std::string(progress_bars, '=').c_str(),
37                std::string(50 - progress_bars, ' ').c_str(),
38                progress_percent);
39    };
40
41    if (data->diff_params->visual_mode) {
42        // Visual mode: clear
43        LOG_INF("\033[2J\033[H");  // Clear screen and move cursor to top-left
44
45        print_progress_bar(step, total_steps);
46
47        LOG_INF("\n");
48
49        std::string current_text = " ";
50
51        for (int32_t i = data->n_input; i < n_tokens; i++) {
52            std::string token_str;
53            if (tokens[i] != llama_vocab_mask(data->vocab)) {
54                char piece[256];
55                int  n_chars = llama_token_to_piece(data->vocab, tokens[i], piece, sizeof(piece), 0, false);
56                if (n_chars > 0) {
57                    piece[n_chars] = '\0';
58                    token_str      = piece;
59                }
60            } else {
61                token_str = " ";
62            }
63
64            current_text += token_str;
65        }
66
67        LOG_INF("%s\n", current_text.c_str());
68    } else {
69        print_progress_bar(step, total_steps);
70    }
71
72    return true;
73}
74
75static std::string format_input_text(const std::string & prompt, const std::string & system_prompt, bool use_chat_template, llama_model * model) {
76    if (!use_chat_template) {
77        return prompt;
78    }
79
80    auto chat_templates = common_chat_templates_init(model, "");
81    common_chat_templates_inputs inputs;
82    common_chat_msg system_msg;
83
84    if (!system_prompt.empty()) {
85        system_msg.role = "system";
86        system_msg.content = system_prompt;
87        inputs.messages.push_back(system_msg);
88    }
89
90    common_chat_msg user_msg;
91    user_msg.role = "user";
92    user_msg.content = prompt;
93
94    inputs.messages.push_back(user_msg);
95    inputs.add_generation_prompt = true;
96
97    auto result = common_chat_templates_apply(chat_templates.get(), inputs);
98
99    return result.prompt;
100}
101
102int main(int argc, char ** argv) {
103    std::setlocale(LC_NUMERIC, "C");
104
105    ggml_time_init();
106
107    common_params params;
108
109    common_init();
110
111    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_DIFFUSION)) {
112        return 1;
113    }
114
115    llama_backend_init();
116
117    llama_model_params model_params = llama_model_default_params();
118    model_params.n_gpu_layers       = params.n_gpu_layers;
119    model_params.devices            = params.devices.data();
120    model_params.load_mode          = params.load_mode;
121    model_params.check_tensors      = params.check_tensors;
122
123    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params);
124    if (!model) {
125        LOG_ERR("error: failed to load model '%s'\n", params.model.path.c_str());
126        return 1;
127    }
128
129    if (!llama_model_is_diffusion(model)) {
130        LOG_ERR("error: unsupported model for diffusion");
131        llama_model_free(model);
132        return 1;
133    }
134
135    llama_context_params ctx_params = llama_context_default_params();
136    ctx_params.n_ctx                = params.n_ctx;
137    ctx_params.n_batch              = params.n_batch;
138    ctx_params.n_ubatch             = params.n_ubatch;
139    ctx_params.flash_attn_type      = params.flash_attn_type;
140    ctx_params.no_perf              = params.no_perf;
141    ctx_params.type_k               = params.cache_type_k;
142    ctx_params.type_v               = params.cache_type_v;
143
144    llama_context * ctx = llama_init_from_model(model, ctx_params);
145    if (!ctx) {
146        LOG_ERR("error: failed to create context\n");
147        llama_model_free(model);
148        return 1;
149    }
150
151    llama_set_n_threads(ctx, params.cpuparams.n_threads, params.cpuparams_batch.n_threads);
152
153    const llama_vocab * vocab            = llama_model_get_vocab(model);
154
155    std::string         formatted_prompt = format_input_text(params.prompt, params.system_prompt, params.enable_chat_template, model);
156
157    std::vector<llama_token> input_tokens = common_tokenize(vocab,
158                                                            formatted_prompt,
159                                                            /*add special tokens*/ true,
160                                                            /*parse special*/ true);
161
162    int n_input = input_tokens.size();
163
164    if (static_cast<uint32_t>(n_input) >= llama_n_ctx(ctx)) {
165        LOG_ERR("error: input too long (%d tokens), max context is %d\n", n_input, llama_n_ctx(ctx));
166        llama_free(ctx);
167        llama_model_free(model);
168        return 1;
169    }
170
171    llama_token mask_token_id = llama_vocab_mask(vocab);
172
173    GGML_ASSERT(mask_token_id != LLAMA_TOKEN_NULL);
174
175    bool visual_mode = params.diffusion.visual_mode;
176
177    int32_t                  n_generated = 0;
178    std::vector<llama_token> output_tokens(params.n_ubatch);
179
180    struct diffusion_params diff_params;
181
182    char shift_logits_str[8];
183    if (llama_model_meta_val_str(model, "diffusion.shift_logits", shift_logits_str, sizeof(shift_logits_str)) >= 0) {
184        diff_params.shift_logits = (strcmp(shift_logits_str, "true") == 0);
185    } else {
186        diff_params.shift_logits = true;
187    }
188
189    //Use either eps or block length, but not both
190    GGML_ASSERT((params.diffusion.eps == 0) ^ (params.diffusion.block_length == 0));
191
192    if (params.diffusion.eps) {
193        diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED;
194        diff_params.eps      = params.diffusion.eps;
195    } else if (params.diffusion.block_length) {
196        diff_params.schedule     = DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED;
197        diff_params.block_length = params.diffusion.block_length;
198    }
199
200    diff_params.mask_token_id    = mask_token_id;
201    diff_params.seed             = params.sampling.seed;
202    diff_params.temperature      = params.sampling.temp;
203    diff_params.steps            = params.diffusion.steps;
204    diff_params.algorithm        = static_cast<diffusion_algorithm>(params.diffusion.algorithm);
205    diff_params.max_length       = params.n_ubatch;
206    diff_params.top_p            = params.sampling.top_p;
207    diff_params.top_k            = params.sampling.top_k;
208    diff_params.visual_mode      = params.diffusion.visual_mode;
209    diff_params.add_gumbel_noise = params.diffusion.add_gumbel_noise;
210
211    diff_params.step_callback           = diffusion_step_callback;
212    callback_data cb_data               = { &diff_params, vocab, n_input };
213    diff_params.step_callback_user_data = &cb_data;
214
215    const char * alg_names[]   = {
216        "DIFFUSION_ALGORITHM_ORIGIN",
217        "DIFFUSION_ALGORITHM_ENTROPY_BASED",
218        "DIFFUSION_ALGORITHM_MARGIN_BASED",
219        "DIFFUSION_ALGORITHM_RANDOM",
220        "DIFFUSION_ALGORITHM_CONFIDENCE_BASED",
221    };
222    const char * sched_names[] = {
223        "DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED",
224        "DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED",
225    };
226    const char * alg_name =
227        (diff_params.algorithm >= 0 && diff_params.algorithm <= 4) ? alg_names[diff_params.algorithm] : "UNKNOWN";
228    const char * sched_name =
229        (diff_params.schedule >= 0 && diff_params.schedule <= 1) ? sched_names[diff_params.schedule] : "UNKNOWN";
230
231    LOG_INF("diffusion_params: - %-25s llama_token      = %d\n", "mask_token_id", mask_token_id);
232    LOG_INF("diffusion_params: - %-25s u32              = %d\n", "steps", diff_params.steps);
233    LOG_INF("diffusion_params: - %-25s u32              = %d\n", "max_length", diff_params.max_length);
234    LOG_INF("diffusion_params: - %-25s enum             = %d (%s)\n", "algorithm", diff_params.algorithm, alg_name);
235    LOG_INF("diffusion_params: - %-25s enum             = %d (%s)\n", "schedule", diff_params.schedule, sched_name);
236    LOG_INF("diffusion_params: - %-25s f32              = %.3f\n", "temperature", diff_params.temperature);
237    if (diff_params.schedule == DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED) {
238        LOG_INF("diffusion_params: - %-25s f32              = %.6f\n", "eps", diff_params.eps);
239        LOG_INF("diffusion_params: - %-25s f32              = %.3f\n", "alg_temp", diff_params.alg_temp);
240    }
241    if (diff_params.schedule == DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED) {
242        LOG_INF("diffusion_params: - %-25s u32              = %d\n", "block_length", diff_params.block_length);
243        LOG_INF("diffusion_params: - %-25s f32              = %.3f\n", "cfg_scale", diff_params.cfg_scale);
244    }
245
246    diffusion_generate(ctx, input_tokens.data(), output_tokens.data(), n_input, diff_params, n_generated);
247
248    if (n_generated > 0) {
249        if (visual_mode) {
250            //clear screen and move cursor to top-left
251            LOG_INF("\033[2J\033[H");
252        }
253
254        output_tokens.erase(output_tokens.begin(), output_tokens.begin() + n_input);
255        std::string output_data = common_detokenize(vocab, output_tokens, false);
256        LOG_INF("\n%s\n", output_data.c_str());
257    } else {
258        LOG_INF("Error: diffusion generation failed\n");
259    }
260
261    llama_free(ctx);
262    llama_model_free(model);
263    llama_backend_free();
264
265    return 0;
266}