infill.cpp 21.2 KB
Newer Older
xuxzh1's avatar
update  
xuxzh1 committed
1
#include "arg.h"
xuxzh1's avatar
init  
xuxzh1 committed
2
3
#include "common.h"
#include "console.h"
xuxzh1's avatar
update  
xuxzh1 committed
4
5
#include "sampling.h"
#include "log.h"
xuxzh1's avatar
init  
xuxzh1 committed
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
#include "llama.h"

#include <cassert>
#include <cinttypes>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
#include <signal.h>
#include <unistd.h>
#elif defined (_WIN32)
#define WIN32_LEAN_AND_MEAN
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#include <signal.h>
#endif

#if defined(_MSC_VER)
#pragma warning(disable: 4244 4267) // possible loss of data
#endif

static llama_context           ** g_ctx;
static llama_model             ** g_model;
xuxzh1's avatar
update  
xuxzh1 committed
38
39
static common_sampler          ** g_smpl;
static common_params            * g_params;
xuxzh1's avatar
init  
xuxzh1 committed
40
41
42
43
44
45
46
47
48
49
50
51
52
static std::vector<llama_token> * g_input_tokens;
static std::ostringstream       * g_output_ss;
static std::vector<llama_token> * g_output_tokens;

static bool is_interacting = false;

#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
static void sigint_handler(int signo) {
    if (signo == SIGINT) {
        if (!is_interacting) {
            is_interacting = true;
        } else {
            console::cleanup();
xuxzh1's avatar
update  
xuxzh1 committed
53
54
55
56
57
58
59
            LOG("\n");
            common_perf_print(*g_ctx, *g_smpl);

            // make sure all logs are flushed
            LOG("Interrupted by user\n");
            common_log_pause(common_log_main());

xuxzh1's avatar
init  
xuxzh1 committed
60
61
62
63
64
65
66
            _exit(130);
        }
    }
}
#endif

int main(int argc, char ** argv) {
xuxzh1's avatar
update  
xuxzh1 committed
67
    common_params params;
xuxzh1's avatar
init  
xuxzh1 committed
68
69
    g_params = &params;

xuxzh1's avatar
update  
xuxzh1 committed
70
    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_INFILL)) {
xuxzh1's avatar
init  
xuxzh1 committed
71
72
73
        return 1;
    }

xuxzh1's avatar
update  
xuxzh1 committed
74
75
76
    common_init();

    auto & sparams = params.sampling;
xuxzh1's avatar
init  
xuxzh1 committed
77
78
79
80
81

    console::init(params.simple_io, params.use_color);
    atexit([]() { console::cleanup(); });

    if (params.logits_all) {
xuxzh1's avatar
update  
xuxzh1 committed
82
83
84
        LOG_ERR("\n************\n");
        LOG_ERR("%s: please use the 'perplexity' tool for perplexity calculations\n", __func__);
        LOG_ERR("************\n\n");
xuxzh1's avatar
init  
xuxzh1 committed
85
86
87
88
89

        return 0;
    }

    if (params.embedding) {
xuxzh1's avatar
update  
xuxzh1 committed
90
91
92
        LOG_ERR("\n************\n");
        LOG_ERR("%s: please use the 'embedding' tool for embedding calculations\n", __func__);
        LOG_ERR("************\n\n");
xuxzh1's avatar
init  
xuxzh1 committed
93
94
95
96
97

        return 0;
    }

    if (params.n_ctx != 0 && params.n_ctx < 8) {
xuxzh1's avatar
update  
xuxzh1 committed
98
        LOG_WRN("%s: minimum context size is 8, using minimum size.\n", __func__);
xuxzh1's avatar
init  
xuxzh1 committed
99
100
        params.n_ctx = 8;
    }
xuxzh1's avatar
update  
xuxzh1 committed
101

xuxzh1's avatar
init  
xuxzh1 committed
102
    if (!params.interactive_first && (params.input_prefix.empty() && params.input_suffix.empty())) {
xuxzh1's avatar
update  
xuxzh1 committed
103
104
105
        LOG_ERR("\n************\n");
        LOG_ERR("%s: please use '--interactive_first' or specify '--in_prefix' and/or '--in_suffix'\n", __func__);
        LOG_ERR("************\n\n");
xuxzh1's avatar
init  
xuxzh1 committed
106
107
108
109
110

        return 0;
    }

    if (params.rope_freq_base != 0.0) {
xuxzh1's avatar
update  
xuxzh1 committed
111
        LOG_WRN("%s: changing RoPE frequency base to %g.\n", __func__, params.rope_freq_base);
xuxzh1's avatar
init  
xuxzh1 committed
112
113
114
    }

    if (params.rope_freq_scale != 0.0) {
xuxzh1's avatar
update  
xuxzh1 committed
115
        LOG_WRN("%s: scaling RoPE frequency by %g.\n", __func__, params.rope_freq_scale);
xuxzh1's avatar
init  
xuxzh1 committed
116
117
    }

xuxzh1's avatar
update  
xuxzh1 committed
118
    LOG_INF("%s: llama backend init\n", __func__);
xuxzh1's avatar
init  
xuxzh1 committed
119
120
121
    llama_backend_init();
    llama_numa_init(params.numa);

xuxzh1's avatar
update  
xuxzh1 committed
122
123
124
    llama_model * model = nullptr;
    llama_context * ctx = nullptr;
    common_sampler * smpl = nullptr;
xuxzh1's avatar
init  
xuxzh1 committed
125
126
127

    g_model = &model;
    g_ctx = &ctx;
xuxzh1's avatar
update  
xuxzh1 committed
128
    g_smpl = &smpl;
xuxzh1's avatar
init  
xuxzh1 committed
129
130

    // load the model and apply lora adapter, if any
xuxzh1's avatar
update  
xuxzh1 committed
131
132
    LOG_INF("%s: load the model and apply lora adapter, if any\n", __func__);
    common_init_result llama_init = common_init_from_params(params);
xuxzh1's avatar
init  
xuxzh1 committed
133
134
135
136
137

    model = llama_init.model;
    ctx = llama_init.context;

    if (model == NULL) {
xuxzh1's avatar
update  
xuxzh1 committed
138
        LOG_ERR("%s: unable to load model\n", __func__);
xuxzh1's avatar
init  
xuxzh1 committed
139
140
141
142
143
        return 1;
    }

    const int n_ctx_train = llama_n_ctx_train(model);
    const int n_ctx = llama_n_ctx(ctx);
xuxzh1's avatar
update  
xuxzh1 committed
144
    LOG_DBG("n_ctx: %d\n", n_ctx);
xuxzh1's avatar
init  
xuxzh1 committed
145
146

    if (n_ctx > n_ctx_train) {
xuxzh1's avatar
update  
xuxzh1 committed
147
        LOG_WRN("%s: model was trained on only %d context tokens (%d specified)\n", __func__, n_ctx_train, n_ctx);
xuxzh1's avatar
init  
xuxzh1 committed
148
149
150
151
    }

    // print system information
    {
xuxzh1's avatar
update  
xuxzh1 committed
152
153
        LOG_INF("\n");
        LOG_INF("%s\n", common_params_get_system_info(params).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
154
    }
xuxzh1's avatar
update  
xuxzh1 committed
155
156
    const bool add_bos = llama_add_bos_token(model);
    GGML_ASSERT(!llama_add_eos_token(model));
xuxzh1's avatar
init  
xuxzh1 committed
157
158
159

    std::vector<llama_token> embd_inp;
    std::vector<llama_token> embd_end;
xuxzh1's avatar
update  
xuxzh1 committed
160
161
    std::vector<llama_token> inp_pfx = common_tokenize(ctx, params.input_prefix, false);
    std::vector<llama_token> inp_sfx = common_tokenize(ctx, params.input_suffix, false);
xuxzh1's avatar
init  
xuxzh1 committed
162

xuxzh1's avatar
update  
xuxzh1 committed
163
164
    GGML_ASSERT(llama_token_fim_pre(model) >= 0);
    GGML_ASSERT(llama_token_fim_suf(model) >= 0);
xuxzh1's avatar
init  
xuxzh1 committed
165

xuxzh1's avatar
update  
xuxzh1 committed
166
167
    inp_pfx.insert(inp_pfx.begin(), llama_token_fim_pre(model));
    inp_sfx.insert(inp_sfx.begin(), llama_token_fim_suf(model));
xuxzh1's avatar
init  
xuxzh1 committed
168
169
170
171
172
173
174
175

    embd_inp = params.spm_infill ? inp_sfx : inp_pfx;
    embd_end = params.spm_infill ? inp_pfx : inp_sfx;
    if (add_bos) {
        embd_inp.insert(embd_inp.begin(), llama_token_bos(model));
    }
    embd_inp.insert(embd_inp.end(), embd_end.begin(), embd_end.end());

xuxzh1's avatar
update  
xuxzh1 committed
176
    const llama_token middle_token = llama_token_fim_mid(model);
xuxzh1's avatar
init  
xuxzh1 committed
177
178
179
180
    if (middle_token >= 0) {
        embd_inp.push_back(middle_token);
    }

xuxzh1's avatar
update  
xuxzh1 committed
181
182
183
184
    LOG_DBG("add_bos: %d\n", add_bos);
    LOG_DBG("prefix: \"%s\"\n", params.input_prefix.c_str());
    LOG_DBG("suffix: \"%s\"\n", params.input_suffix.c_str());
    LOG_DBG("tokens: %s\n", string_from(ctx, embd_inp).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
185
186
187
188

    // Should not run without any tokens
    if (embd_inp.empty()) {
        embd_inp.push_back(llama_token_bos(model));
xuxzh1's avatar
update  
xuxzh1 committed
189
        LOG_WRN("embd_inp was considered empty and bos was added: %s\n", string_from(ctx, embd_inp).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
190
191
192
    }

    if ((int) embd_inp.size() > n_ctx - 4) {
xuxzh1's avatar
update  
xuxzh1 committed
193
        LOG_ERR("%s: prompt is too long (%d tokens, max %d)\n", __func__, (int) embd_inp.size(), n_ctx - 4);
xuxzh1's avatar
init  
xuxzh1 committed
194
195
196
197
198
199
200
201
        return 1;
    }

    // number of tokens to keep when resetting context
    if (params.n_keep < 0 || params.n_keep > (int) embd_inp.size()) {
        params.n_keep = (int)embd_inp.size();
    }

xuxzh1's avatar
update  
xuxzh1 committed
202
203
    LOG_INF("inp_pfx: %s\n", string_from(ctx, inp_pfx).c_str());
    LOG_INF("inp_sfx: %s\n", string_from(ctx, inp_sfx).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
204
205
206
207
208
209
210

    // enable interactive mode if interactive start is specified
    if (params.interactive_first) {
        params.interactive = true;
    }

    if (params.verbose_prompt) {
xuxzh1's avatar
update  
xuxzh1 committed
211
212
213
        LOG_INF("\n");
        LOG_INF("%s: prompt: '%s'\n", __func__, params.prompt.c_str());
        LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, embd_inp.size());
xuxzh1's avatar
init  
xuxzh1 committed
214
        for (int i = 0; i < (int) embd_inp.size(); i++) {
xuxzh1's avatar
update  
xuxzh1 committed
215
            LOG_INF("%6d -> '%s'\n", embd_inp[i], common_token_to_piece(ctx, embd_inp[i]).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
216
217
218
        }

        if (params.n_keep > 0) {
xuxzh1's avatar
update  
xuxzh1 committed
219
        LOG_INF("%s: static prompt based on n_keep: '", __func__);
xuxzh1's avatar
init  
xuxzh1 committed
220
            for (int i = 0; i < params.n_keep; i++) {
xuxzh1's avatar
update  
xuxzh1 committed
221
                LOG_CNT("%s", common_token_to_piece(ctx, embd_inp[i]).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
222
            }
xuxzh1's avatar
update  
xuxzh1 committed
223
            LOG_CNT("'\n");
xuxzh1's avatar
init  
xuxzh1 committed
224
        }
xuxzh1's avatar
update  
xuxzh1 committed
225
        LOG_INF("\n");
xuxzh1's avatar
init  
xuxzh1 committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    }

    if (params.interactive) {
#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
        struct sigaction sigint_action;
        sigint_action.sa_handler = sigint_handler;
        sigemptyset (&sigint_action.sa_mask);
        sigint_action.sa_flags = 0;
        sigaction(SIGINT, &sigint_action, NULL);
#elif defined (_WIN32)
        auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
            return (ctrl_type == CTRL_C_EVENT) ? (sigint_handler(SIGINT), true) : false;
        };
        SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
#endif

xuxzh1's avatar
update  
xuxzh1 committed
242
        LOG_INF("%s: interactive mode on.\n", __func__);
xuxzh1's avatar
init  
xuxzh1 committed
243
244

        if (params.input_prefix_bos) {
xuxzh1's avatar
update  
xuxzh1 committed
245
            LOG_INF("Input prefix with BOS\n");
xuxzh1's avatar
init  
xuxzh1 committed
246
247
248
        }

        if (!params.input_prefix.empty()) {
xuxzh1's avatar
update  
xuxzh1 committed
249
            LOG_INF("Input prefix: '%s'\n", params.input_prefix.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
250
251
252
        }

        if (!params.input_suffix.empty()) {
xuxzh1's avatar
update  
xuxzh1 committed
253
            LOG_INF("Input suffix: '%s'\n", params.input_suffix.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
254
255
        }
    }
xuxzh1's avatar
update  
xuxzh1 committed
256
257
258
259
260
261
262
263
264
265
    smpl = common_sampler_init(model, sparams);

    LOG_INF("sampler seed: %u\n",     common_sampler_get_seed(smpl));
    LOG_INF("sampler params: \n%s\n", sparams.print().c_str());
    LOG_INF("sampler chain: %s\n",    common_sampler_print(smpl).c_str());

    LOG_INF("generate: n_ctx = %d, n_batch = %d, n_predict = %d, n_keep = %d\n", n_ctx, params.n_batch, params.n_predict, params.n_keep);

    LOG_INF("\n");
    LOG_INF("\n#####  Infill mode  #####\n\n");
xuxzh1's avatar
init  
xuxzh1 committed
266
267
268
269
270
271
272
273
274
275
    if (params.interactive) {
        const char *control_message;
        if (params.multiline_input) {
            control_message = " - To return control to LLaMA, end your input with '\\'.\n"
                              " - To return control without starting a new line, end your input with '/'.\n";
        } else {
            control_message = " - Press Return to return control to LLaMA.\n"
                              " - To return control without starting a new line, end your input with '/'.\n"
                              " - If you want to submit another line, end your input with '\\'.\n";
        }
xuxzh1's avatar
update  
xuxzh1 committed
276
        LOG_INF("== Running in interactive mode. ==\n");
xuxzh1's avatar
init  
xuxzh1 committed
277
#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
xuxzh1's avatar
update  
xuxzh1 committed
278
        LOG_INF(       " - Press Ctrl+C to interject at any time.\n");
xuxzh1's avatar
init  
xuxzh1 committed
279
#endif
xuxzh1's avatar
update  
xuxzh1 committed
280
        LOG_INF(       "%s\n", control_message);
xuxzh1's avatar
init  
xuxzh1 committed
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312

        is_interacting = params.interactive_first;
    }

    bool input_echo = true;

    int n_past     = 0;
    int n_remain   = params.n_predict;
    int n_consumed = 0;

    std::vector<int>   input_tokens;  g_input_tokens  = &input_tokens;
    std::vector<int>   output_tokens; g_output_tokens = &output_tokens;
    std::ostringstream output_ss;     g_output_ss     = &output_ss;

    // the first thing we will do is to output the prompt, so set color accordingly
    console::set_display(console::prompt);

    std::vector<llama_token> embd;

    while (n_remain != 0 || params.interactive) {
        // predict
        if (!embd.empty()) {
            // Note: n_ctx - 4 here is to match the logic for commandline prompt handling via
            // --prompt or --file which uses the same value.
            int max_embd_size = n_ctx - 4;

            // Ensure the input doesn't exceed the context size by truncating embd if necessary.
            if ((int) embd.size() > max_embd_size) {
                const int skipped_tokens = (int) embd.size() - max_embd_size;
                embd.resize(max_embd_size);

                console::set_display(console::error);
xuxzh1's avatar
update  
xuxzh1 committed
313
                LOG_WRN("<<input too long: skipped %d token%s>>", skipped_tokens, skipped_tokens != 1 ? "s" : "");
xuxzh1's avatar
init  
xuxzh1 committed
314
315
316
317
318
319
320
321
322
                console::set_display(console::reset);
            }

            // infinite text generation via context swapping
            // if we run out of context:
            // - take the n_keep first tokens from the original prompt (via n_past)
            // - take half of the last (n_ctx - n_keep) tokens and recompute the logits in batches
            if (n_past + (int) embd.size() > n_ctx) {
                if (params.n_predict == -2) {
xuxzh1's avatar
update  
xuxzh1 committed
323
                    LOG_DBG("\n\n%s: context full and n_predict == -%d => stopping\n", __func__, params.n_predict);
xuxzh1's avatar
init  
xuxzh1 committed
324
325
326
327
328
329
                    break;
                }

                const int n_left    = n_past - params.n_keep - 1;
                const int n_discard = n_left/2;

xuxzh1's avatar
update  
xuxzh1 committed
330
                LOG_DBG("context full, swapping: n_past = %d, n_left = %d, n_ctx = %d, n_keep = %d, n_discard = %d\n",
xuxzh1's avatar
init  
xuxzh1 committed
331
332
333
334
335
336
337
                    n_past, n_left, n_ctx, params.n_keep, n_discard);

                llama_kv_cache_seq_rm (ctx, 0, params.n_keep + 1            , params.n_keep + n_discard + 1);
                llama_kv_cache_seq_add(ctx, 0, params.n_keep + 1 + n_discard, n_past, -n_discard);

                n_past -= n_discard;

xuxzh1's avatar
update  
xuxzh1 committed
338
                LOG_DBG("after swap: n_past = %d\n", n_past);
xuxzh1's avatar
init  
xuxzh1 committed
339

xuxzh1's avatar
update  
xuxzh1 committed
340
                LOG_DBG("embd: %s\n", string_from(ctx, embd).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
341
342
343
344
345
346
347
348
349
350
351

            }

            // evaluate tokens in batches
            // embd is typically prepared beforehand to fit within a batch, but not always
            for (int i = 0; i < (int) embd.size(); i += params.n_batch) {
                int n_eval = (int) embd.size() - i;
                if (n_eval > params.n_batch) {
                    n_eval = params.n_batch;
                }

xuxzh1's avatar
update  
xuxzh1 committed
352
                LOG_DBG("eval: %s\n", string_from(ctx, embd).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
353

xuxzh1's avatar
update  
xuxzh1 committed
354
355
                if (llama_decode(ctx, llama_batch_get_one(&embd[i], n_eval))) {
                    LOG_ERR("%s : failed to eval\n", __func__);
xuxzh1's avatar
init  
xuxzh1 committed
356
357
358
359
360
                    return 1;
                }

                n_past += n_eval;

xuxzh1's avatar
update  
xuxzh1 committed
361
                LOG_DBG("n_past = %d\n", n_past);
xuxzh1's avatar
init  
xuxzh1 committed
362
363
364
365
366
367
368
            }

        }

        embd.clear();

        if ((int) embd_inp.size() <= n_consumed && !is_interacting) {
xuxzh1's avatar
update  
xuxzh1 committed
369
            const llama_token id = common_sampler_sample(smpl, ctx, -1);
xuxzh1's avatar
init  
xuxzh1 committed
370

xuxzh1's avatar
update  
xuxzh1 committed
371
            common_sampler_accept(smpl, id, true);
xuxzh1's avatar
init  
xuxzh1 committed
372

xuxzh1's avatar
update  
xuxzh1 committed
373
            // LOG_DBG("last: %s\n", string_from(ctx, smpl->prev.to_vector()).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
374
375
376
377
378
379
380
381
382

            embd.push_back(id);

            // echo this to console
            input_echo = true;

            // decrement remaining sampling budget
            --n_remain;

xuxzh1's avatar
update  
xuxzh1 committed
383
            LOG_DBG("n_remain: %d\n", n_remain);
xuxzh1's avatar
init  
xuxzh1 committed
384
385
        } else {
            // some user input remains from prompt or interaction, forward it to processing
xuxzh1's avatar
update  
xuxzh1 committed
386
            LOG_DBG("embd_inp.size(): %d, n_consumed: %d\n", (int) embd_inp.size(), n_consumed);
xuxzh1's avatar
init  
xuxzh1 committed
387
388
389
390
391
            while ((int) embd_inp.size() > n_consumed) {
                embd.push_back(embd_inp[n_consumed]);

                // push the prompt in the sampling context in order to apply repetition penalties later
                // for the prompt, we don't apply grammar rules
xuxzh1's avatar
update  
xuxzh1 committed
392
                common_sampler_accept(smpl, embd_inp[n_consumed], false);
xuxzh1's avatar
init  
xuxzh1 committed
393
394
395
396
397
398
399
400
401
402
403

                ++n_consumed;
                if ((int) embd.size() >= params.n_batch) {
                    break;
                }
            }
        }

        // display text
        if (input_echo) {
            for (auto id : embd) {
xuxzh1's avatar
update  
xuxzh1 committed
404
405
                const std::string token_str = common_token_to_piece(ctx, id);
                LOG("%s", token_str.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422

                if (embd.size() > 1) {
                    input_tokens.push_back(id);
                } else {
                    output_tokens.push_back(id);
                    output_ss << token_str;
                }
            }
        }
        // reset color to default if we there is no pending user input
        if (input_echo && (int) embd_inp.size() == n_consumed) {
            console::set_display(console::reset);
        }

        // if not currently processing queued inputs;
        if ((int) embd_inp.size() <= n_consumed) {
            // deal with eot token in infill mode
xuxzh1's avatar
update  
xuxzh1 committed
423
            if ((common_sampler_last(smpl) == llama_token_eot(model) || is_interacting) && params.interactive){
xuxzh1's avatar
init  
xuxzh1 committed
424
425
                if (is_interacting && !params.interactive_first) {
                    // print an eot token
xuxzh1's avatar
update  
xuxzh1 committed
426
                    LOG("%s", common_token_to_piece(ctx, llama_token_eot(model)).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
427
                }
xuxzh1's avatar
update  
xuxzh1 committed
428
                LOG("\n");
xuxzh1's avatar
init  
xuxzh1 committed
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
                console::set_display(console::user_input);
                std::string buffer;
                std::string line;
                bool another_line=true;
                // set a new prefix via stdin
                do {
                    another_line = console::readline(line, params.multiline_input);
                    buffer += line;
                } while (another_line);
                // check if we got an empty line, if so we use the old input
                if (!buffer.empty() && !(buffer.length() == 1 && buffer[0] == '\n')) {
                    params.input_prefix = buffer;
                }
                buffer.clear();
                // set a new suffix via stdin
                do {
                    another_line = console::readline(line, params.multiline_input);
                    buffer += line;
                } while (another_line);
                // check if we got an empty line
                if (!buffer.empty() && !(buffer.length() == 1 && buffer[0] == '\n')) {
                    params.input_suffix = buffer;
                }
                buffer.clear();
                // done taking input, reset color
                console::set_display(console::reset);

                if (params.escape) {
                    //process escape sequences, for the initial prompt this is done in common.cpp when we load the params, but for the interactive mode we need to do it here
                    string_process_escapes(params.input_prefix);
                    string_process_escapes(params.input_suffix);
                }

                // tokenize new prefix and suffix
xuxzh1's avatar
update  
xuxzh1 committed
463
464
                std::vector<llama_token> inp_pfx = common_tokenize(ctx, params.input_prefix, false);
                std::vector<llama_token> inp_sfx = common_tokenize(ctx, params.input_suffix, false);
xuxzh1's avatar
init  
xuxzh1 committed
465

xuxzh1's avatar
update  
xuxzh1 committed
466
467
                inp_pfx.insert(inp_pfx.begin(), llama_token_fim_pre(model));
                inp_sfx.insert(inp_sfx.begin(), llama_token_fim_suf(model));
xuxzh1's avatar
init  
xuxzh1 committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486

                embd_inp = params.spm_infill ? inp_sfx : inp_pfx;
                embd_end = params.spm_infill ? inp_pfx : inp_sfx;
                if (add_bos) {
                    embd_inp.insert(embd_inp.begin(), llama_token_bos(model));
                }
                embd_inp.insert(embd_inp.end(), embd_end.begin(), embd_end.end());

                if (middle_token >= 0) {
                    embd_inp.push_back(middle_token);
                }

                embd.clear();
                n_remain = params.n_predict;
                n_past = 0;
                n_consumed = 0;
                is_interacting = false;
            }
            // deal with end of generation tokens in interactive mode
xuxzh1's avatar
update  
xuxzh1 committed
487
488
            else if (llama_token_is_eog(model, common_sampler_last(smpl))) {
                LOG_DBG("found EOS token\n");
xuxzh1's avatar
init  
xuxzh1 committed
489
490
491
492

                if (params.interactive) {

                    is_interacting = true;
xuxzh1's avatar
update  
xuxzh1 committed
493
                    LOG("\n");
xuxzh1's avatar
init  
xuxzh1 committed
494
495
496
497
498
                    console::set_display(console::user_input);
               }
            }

            if (n_past > 0 && is_interacting && !params.interactive) {
xuxzh1's avatar
update  
xuxzh1 committed
499
                LOG_DBG("waiting for user input\n");
xuxzh1's avatar
init  
xuxzh1 committed
500
501

                if (params.input_prefix_bos) {
xuxzh1's avatar
update  
xuxzh1 committed
502
                    LOG_DBG("adding input prefix BOS token\n");
xuxzh1's avatar
init  
xuxzh1 committed
503
504
505
506
507
                    embd_inp.push_back(llama_token_bos(model));
                }

                std::string buffer;
                if (!params.input_prefix.empty()) {
xuxzh1's avatar
update  
xuxzh1 committed
508
                    LOG_DBG("appending input prefix: '%s'\n", params.input_prefix.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
509
                    buffer += params.input_prefix;
xuxzh1's avatar
update  
xuxzh1 committed
510
                    LOG("%s", buffer.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
                }

                std::string line;
                bool another_line = true;
                do {
                    another_line = console::readline(line, params.multiline_input);
                    buffer += line;
                } while (another_line);

                // done taking input, reset color
                console::set_display(console::reset);

                // Add tokens to embd only if the input buffer is non-empty
                // Entering a empty line lets the user pass control back
                if (buffer.length() > 1) {
                    // append input suffix if any
                    if (!params.input_suffix.empty()) {
xuxzh1's avatar
update  
xuxzh1 committed
528
                        LOG_DBG("appending input suffix: '%s'\n", params.input_suffix.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
529
                        buffer += params.input_suffix;
xuxzh1's avatar
update  
xuxzh1 committed
530
                        LOG("%s", params.input_suffix.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
531
532
                    }

xuxzh1's avatar
update  
xuxzh1 committed
533
                    LOG_DBG("buffer: '%s'\n", buffer.c_str());
xuxzh1's avatar
init  
xuxzh1 committed
534
535
536

                    const size_t original_size = embd_inp.size();

xuxzh1's avatar
update  
xuxzh1 committed
537
538
                    const auto line_inp = common_tokenize(ctx, buffer, false);
                    LOG_DBG("input tokens: %s\n", string_from(ctx, line_inp).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
539
540
541
542
543
544

                    embd_inp.insert(embd_inp.end(), line_inp.begin(), line_inp.end());

                    for (size_t i = original_size; i < embd_inp.size(); ++i) {
                        const llama_token token = embd_inp[i];
                        output_tokens.push_back(token);
xuxzh1's avatar
update  
xuxzh1 committed
545
                        output_ss << common_token_to_piece(ctx, token);
xuxzh1's avatar
init  
xuxzh1 committed
546
547
548
                    }

                    n_remain -= line_inp.size();
xuxzh1's avatar
update  
xuxzh1 committed
549
                    LOG_DBG("n_remain: %d\n", n_remain);
xuxzh1's avatar
init  
xuxzh1 committed
550
                } else {
xuxzh1's avatar
update  
xuxzh1 committed
551
                    LOG_DBG("empty line, passing control back\n");
xuxzh1's avatar
init  
xuxzh1 committed
552
553
554
555
556
557
558
                }

                input_echo = false; // do not echo this again
            }

            if (n_past > 0) {
                if (is_interacting) {
xuxzh1's avatar
update  
xuxzh1 committed
559
                    common_sampler_reset(smpl);
xuxzh1's avatar
init  
xuxzh1 committed
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
                }
                is_interacting = false;
            }
        }

        // end of generation
        if (!embd.empty() && llama_token_is_eog(model, embd.back()) && !params.interactive) {
            break;
        }

        // In interactive mode, respect the maximum number of tokens and drop back to user input when reached.
        // We skip this logic when n_predict == -1 (infinite) or -2 (stop at context size).
        if (params.interactive && n_remain <= 0 && params.n_predict >= 0) {
            n_remain = params.n_predict;
            is_interacting = true;
        }
    }
    if (!params.interactive && n_remain <= 0) {
xuxzh1's avatar
update  
xuxzh1 committed
578
        LOG("%s", common_token_to_piece(ctx, llama_token_eot(model)).c_str());
xuxzh1's avatar
init  
xuxzh1 committed
579
580
    }

xuxzh1's avatar
update  
xuxzh1 committed
581
582
    LOG("\n");
    common_perf_print(ctx, smpl);
xuxzh1's avatar
init  
xuxzh1 committed
583
584
585
586

    llama_free(ctx);
    llama_free_model(model);

xuxzh1's avatar
update  
xuxzh1 committed
587
    common_sampler_free(smpl);
xuxzh1's avatar
init  
xuxzh1 committed
588
589
590
591
    llama_backend_free();

    return 0;
}