llama.go 10.9 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
package llama

Michael Yang's avatar
Michael Yang committed
3
/*
Michael Yang's avatar
Michael Yang committed
4
#cgo CPPFLAGS: -O3 -Wall -Wextra -Wno-unused-function -Wno-unused-variable -DNDEBUG -DGGML_USE_K_QUANTS
Michael Yang's avatar
Michael Yang committed
5
#cgo CXXFLAGS: -std=gnu++11
6
7
#cgo darwin CPPFLAGS:  -DGGML_USE_ACCELERATE
#cgo darwin,arm64 CPPFLAGS: -DGGML_USE_METAL -DGGML_METAL_NDEBUG
Michael Yang's avatar
Michael Yang committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
#cgo darwin LDFLAGS: -framework Accelerate -framework Foundation -framework Metal -framework MetalKit -framework MetalPerformanceShaders
#include <stdlib.h>
#include "llama.h"

struct llama_sample_options
{
	float repeat_penalty;
	float frequency_penalty;
	float presence_penalty;
	float temperature;
	int32_t top_k;
	float top_p;
	float tfs_z;
	float typical_p;
	int mirostat;
	float mirostat_tau;
	float mirostat_eta;
Michael Yang's avatar
Michael Yang committed
25
	bool penalize_newline;
Michael Yang's avatar
Michael Yang committed
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
};

llama_token llama_sample(
		struct llama_context *ctx,
		struct llama_token_data *candidates,
		size_t n_candidates,
		const llama_token *last_tokens,
		size_t n_last_tokens,
		struct llama_sample_options *opts)
{
	llama_token_data_array candidates_p = {
		candidates,
		n_candidates,
		false,
	};

Michael Yang's avatar
Michael Yang committed
42
43
	struct llama_token_data newline = candidates_p.data[llama_token_nl()];

Michael Yang's avatar
Michael Yang committed
44
45
46
47
48
49
50
51
52
53
	llama_sample_repetition_penalty(
		ctx, &candidates_p,
		last_tokens, n_last_tokens,
		opts->repeat_penalty);

	llama_sample_frequency_and_presence_penalties(
		ctx, &candidates_p,
		last_tokens, n_last_tokens,
		opts->frequency_penalty, opts->presence_penalty);

Michael Yang's avatar
Michael Yang committed
54
55
56
57
	if (!opts->penalize_newline) {
		candidates_p.data[llama_token_nl()] = newline;
	}

Michael Yang's avatar
Michael Yang committed
58
59
60
	if (opts->temperature <= 0) {
		return llama_sample_token_greedy(ctx, &candidates_p);
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
61

Michael Yang's avatar
Michael Yang committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
	if (opts->mirostat == 1) {
		int mirostat_m = 100;
		float mirostat_mu = 2.0f * opts->mirostat_tau;
		llama_sample_temperature(ctx, &candidates_p, opts->temperature);
		return llama_sample_token_mirostat(
			ctx, &candidates_p,
			opts->mirostat_tau, opts->mirostat_eta,
			mirostat_m, &mirostat_mu);
	} else if (opts->mirostat == 2) {
		float mirostat_mu = 2.0f * opts->mirostat_tau;
		llama_sample_temperature(ctx, &candidates_p, opts->temperature);
		return llama_sample_token_mirostat_v2(
			ctx, &candidates_p,
			opts->mirostat_tau, opts->mirostat_eta,
			&mirostat_mu);
	} else {
		llama_sample_top_k(ctx, &candidates_p, opts->top_k, 1);
		llama_sample_tail_free(ctx, &candidates_p, opts->tfs_z, 1);
		llama_sample_typical(ctx, &candidates_p, opts->typical_p, 1);
		llama_sample_top_p(ctx, &candidates_p, opts->top_p, 1);
		llama_sample_temperature(ctx, &candidates_p, opts->temperature);
		return llama_sample_token(ctx, &candidates_p);
	}
}
*/
import "C"
88

Jeffrey Morgan's avatar
Jeffrey Morgan committed
89
import (
Michael Yang's avatar
Michael Yang committed
90
	"bytes"
Michael Yang's avatar
Michael Yang committed
91
	"embed"
Michael Yang's avatar
Michael Yang committed
92
	"errors"
93
	"fmt"
Michael Yang's avatar
Michael Yang committed
94
	"io"
Michael Yang's avatar
Michael Yang committed
95
	"log"
Michael Yang's avatar
Michael Yang committed
96
	"os"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
97
	"strings"
Michael Yang's avatar
Michael Yang committed
98
	"sync"
Michael Yang's avatar
Michael Yang committed
99
	"unicode/utf8"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
100
	"unsafe"
Michael Yang's avatar
Michael Yang committed
101
102

	"github.com/jmorganca/ollama/api"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
103
104
)

Michael Yang's avatar
Michael Yang committed
105
106
107
//go:embed ggml-metal.metal
var fs embed.FS

Michael Yang's avatar
Michael Yang committed
108
type LLM struct {
Michael Yang's avatar
Michael Yang committed
109
110
111
	params *C.struct_llama_context_params
	model  *C.struct_llama_model
	ctx    *C.struct_llama_context
Jeffrey Morgan's avatar
Jeffrey Morgan committed
112

Michael Yang's avatar
Michael Yang committed
113
114
115
116
	last   []C.llama_token
	embd   []C.llama_token
	cursor int

Michael Yang's avatar
Michael Yang committed
117
118
119
	mu sync.Mutex
	gc bool

Michael Yang's avatar
Michael Yang committed
120
121
	api.Options
}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
122

Michael Yang's avatar
Michael Yang committed
123
func New(model string, opts api.Options) (*LLM, error) {
Michael Yang's avatar
Michael Yang committed
124
125
	if _, err := os.Stat(model); err != nil {
		return nil, err
Jeffrey Morgan's avatar
Jeffrey Morgan committed
126
127
	}

Michael Yang's avatar
Michael Yang committed
128
	llm := LLM{Options: opts}
Michael Yang's avatar
Michael Yang committed
129

Michael Yang's avatar
Michael Yang committed
130
	C.llama_backend_init(C.bool(llm.UseNUMA))
Michael Yang's avatar
Michael Yang committed
131
132
133
134
135

	params := C.llama_context_default_params()
	params.seed = C.uint(llm.Seed)
	params.n_ctx = C.int(llm.NumCtx)
	params.n_batch = C.int(llm.NumBatch)
Michael Yang's avatar
Michael Yang committed
136
	params.n_gqa = C.int(llm.NumGQA)
Michael Yang's avatar
Michael Yang committed
137
138
139
140
141
142
143
144
145
	params.n_gpu_layers = C.int(llm.NumGPU)
	params.main_gpu = C.int(llm.MainGPU)
	params.low_vram = C.bool(llm.LowVRAM)
	params.f16_kv = C.bool(llm.F16KV)
	params.logits_all = C.bool(llm.LogitsAll)
	params.vocab_only = C.bool(llm.VocabOnly)
	params.use_mmap = C.bool(llm.UseMMap)
	params.use_mlock = C.bool(llm.UseMLock)
	params.embedding = C.bool(llm.EmbeddingOnly)
146
147
	params.rope_freq_base = C.float(llm.RopeFrequencyBase)
	params.rope_freq_scale = C.float(llm.RopeFrequencyScale)
Michael Yang's avatar
Michael Yang committed
148
149
150
151
152
153
	llm.params = &params

	cModel := C.CString(model)
	defer C.free(unsafe.Pointer(cModel))

	llm.model = C.llama_load_model_from_file(cModel, params)
154
155
156
157
	if llm.model == nil {
		return nil, errors.New("failed to load model")
	}

Michael Yang's avatar
Michael Yang committed
158
	llm.ctx = C.llama_new_context_with_model(llm.model, params)
159
160
161
	if llm.ctx == nil {
		return nil, errors.New("failed to create context")
	}
Michael Yang's avatar
Michael Yang committed
162
163
164
165
166
167
168

	// warm up the model
	bos := []C.llama_token{C.llama_token_bos()}
	C.llama_eval(llm.ctx, unsafe.SliceData(bos), C.int(len(bos)), 0, C.int(opts.NumThread))
	C.llama_reset_timings(llm.ctx)

	return &llm, nil
Jeffrey Morgan's avatar
Jeffrey Morgan committed
169
170
}

Michael Yang's avatar
Michael Yang committed
171
func (llm *LLM) Close() {
Michael Yang's avatar
Michael Yang committed
172
173
174
175
176
	llm.gc = true

	llm.mu.Lock()
	defer llm.mu.Unlock()

Michael Yang's avatar
Michael Yang committed
177
178
179
180
	defer C.llama_free_model(llm.model)
	defer C.llama_free(llm.ctx)

	C.llama_print_timings(llm.ctx)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
181
182
}

Michael Yang's avatar
Michael Yang committed
183
184
var errNeedMoreData = errors.New("need more data")

Michael Yang's avatar
Michael Yang committed
185
func (llm *LLM) Predict(ctx []int, prompt string, fn func(api.GenerateResponse)) error {
Michael Yang's avatar
Michael Yang committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
	C.llama_reset_timings(llm.ctx)

	tokens := make([]C.llama_token, len(ctx))
	for i := range tokens {
		tokens[i] = C.llama_token(ctx[i])
	}

	llm.marshalPrompt(tokens, prompt)

	C.llama_set_rng_seed(llm.ctx, C.uint(llm.Seed))

	var b bytes.Buffer
	for {
		token, err := llm.next()
200
201
202
		if llm.gc {
			return nil
		} else if errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
203
204
205
206
207
			break
		} else if err != nil {
			return err
		}

208
		b.WriteString(llm.Decode(token))
Michael Yang's avatar
Michael Yang committed
209
210
211
212
213
214
215
216
217
218
219

		if err := llm.checkStopConditions(b); err != nil {
			if errors.Is(err, io.EOF) {
				break
			} else if errors.Is(err, errNeedMoreData) {
				continue
			}

			return err
		}

Michael Yang's avatar
Michael Yang committed
220
221
222
		if utf8.Valid(b.Bytes()) || b.Len() >= utf8.UTFMax {
			fn(api.GenerateResponse{Response: b.String()})
			b.Reset()
Michael Yang's avatar
Michael Yang committed
223
		}
Michael Yang's avatar
Michael Yang committed
224
	}
Michael Yang's avatar
Michael Yang committed
225

226
227
228
	embd := make([]int, len(llm.embd))
	for i := range llm.embd {
		embd[i] = int(llm.embd[i])
Jeffrey Morgan's avatar
Jeffrey Morgan committed
229
230
	}

Michael Yang's avatar
Michael Yang committed
231
232
233
	timings := C.llama_get_timings(llm.ctx)
	fn(api.GenerateResponse{
		Done:               true,
234
		Context:            embd,
Michael Yang's avatar
Michael Yang committed
235
236
		SampleCount:        int(timings.n_sample),
		SampleDuration:     parseDurationMs(float64(timings.t_sample_ms)),
Michael Yang's avatar
Michael Yang committed
237
238
239
240
241
242
243
244
245
		PromptEvalCount:    int(timings.n_p_eval),
		PromptEvalDuration: parseDurationMs(float64(timings.t_p_eval_ms)),
		EvalCount:          int(timings.n_eval),
		EvalDuration:       parseDurationMs(float64(timings.t_eval_ms)),
	})

	return nil
}

Michael Yang's avatar
Michael Yang committed
246
func (llm *LLM) checkStopConditions(b bytes.Buffer) error {
247
	for _, stopCondition := range llm.Stop {
248
		if stopCondition == strings.TrimSpace(b.String()) {
Michael Yang's avatar
Michael Yang committed
249
			return io.EOF
250
		} else if strings.HasPrefix(stopCondition, strings.TrimSpace(b.String())) {
Michael Yang's avatar
Michael Yang committed
251
252
253
254
255
256
257
			return errNeedMoreData
		}
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
258
func (llm *LLM) marshalPrompt(ctx []C.llama_token, prompt string) []C.llama_token {
259
	tokens := append(ctx, llm.Encode(prompt)...)
Michael Yang's avatar
Michael Yang committed
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
	if llm.NumKeep < 0 {
		llm.NumKeep = len(tokens)
	}

	// min(llm.NumCtx - 4, llm.NumKeep)
	if llm.NumCtx-4 < llm.NumKeep {
		llm.NumKeep = llm.NumCtx - 4
	}

	if len(tokens) >= llm.NumCtx {
		// truncate input
		numLeft := (llm.NumCtx - llm.NumKeep) / 2
		truncated := tokens[:llm.NumKeep]
		erasedBlocks := (len(tokens) - llm.NumKeep - numLeft - 1) / numLeft
		truncated = append(truncated, tokens[llm.NumKeep+erasedBlocks*numLeft:]...)
		copy(llm.last, tokens[len(tokens)-llm.NumCtx:])

		tokens = truncated
		log.Printf("input truncated: num_ctx=%d num_keep=%d num_left=%d num_tokens=%d", llm.NumCtx, llm.NumKeep, numLeft, len(truncated))
	} else {
		llm.last = make([]C.llama_token, llm.NumCtx-len(tokens))
		llm.last = append(llm.last, tokens...)
	}

	var i int
	for i = 0; i < len(llm.embd) && i < len(tokens) && llm.embd[i] == tokens[i]; i++ {
		// noop
	}

	llm.embd = tokens
	if i == len(tokens) {
		// evaluate at least one token to generate logits
		i--
	}

	llm.cursor = i

	log.Printf("prompt: num_past=%d cached=%v eval=%v", i, len(llm.embd[:i]), len(llm.embd[i:]))
	return tokens
Michael Yang's avatar
Michael Yang committed
299
}
Michael Yang's avatar
Michael Yang committed
300

301
func (llm *LLM) Encode(prompt string) []C.llama_token {
Michael Yang's avatar
Michael Yang committed
302
303
	cPrompt := C.CString(prompt)
	defer C.free(unsafe.Pointer(cPrompt))
Michael Yang's avatar
Michael Yang committed
304

305
	tokens := make([]C.llama_token, len(prompt)+1)
Michael Yang's avatar
Michael Yang committed
306
307
	if n := C.llama_tokenize(llm.ctx, cPrompt, unsafe.SliceData(tokens), C.int(len(tokens)), true); n > 0 {
		return tokens[:n]
Jeffrey Morgan's avatar
Jeffrey Morgan committed
308
309
310
311
312
	}

	return nil
}

313
func (llm *LLM) Decode(tokens ...C.llama_token) string {
Michael Yang's avatar
Michael Yang committed
314
315
316
	var sb strings.Builder
	for _, token := range tokens {
		sb.WriteString(C.GoString(C.llama_token_to_str(llm.ctx, token)))
Jeffrey Morgan's avatar
Jeffrey Morgan committed
317
318
	}

Michael Yang's avatar
Michael Yang committed
319
	return sb.String()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
320
321
}

Michael Yang's avatar
Michael Yang committed
322
func (llm *LLM) next() (C.llama_token, error) {
323
324
325
	llm.mu.Lock()
	defer llm.mu.Unlock()

Michael Yang's avatar
Michael Yang committed
326
327
328
329
	if len(llm.embd) >= llm.NumCtx {
		numLeft := (llm.NumCtx - llm.NumKeep) / 2
		truncated := llm.embd[:llm.NumKeep]
		truncated = append(truncated, llm.embd[len(llm.embd)-numLeft:]...)
Michael Yang's avatar
Michael Yang committed
330

Michael Yang's avatar
Michael Yang committed
331
332
333
334
		llm.embd = truncated
		llm.cursor = llm.NumKeep
		log.Printf("input truncated: num_ctx=%d num_keep=%d num_left=%d num_tokens=%d cursor=%d", llm.NumCtx, llm.NumKeep, numLeft, len(truncated), llm.cursor)
	}
Michael Yang's avatar
Michael Yang committed
335

Michael Yang's avatar
Michael Yang committed
336
	for {
337
338
339
340
		if llm.gc {
			return 0, io.EOF
		}

Michael Yang's avatar
Michael Yang committed
341
		if llm.cursor >= len(llm.embd) {
342
			break
Michael Yang's avatar
Michael Yang committed
343
344
		}

Michael Yang's avatar
Michael Yang committed
345
346
347
		numEval := len(llm.embd) - llm.cursor
		if numEval > llm.NumBatch {
			numEval = llm.NumBatch
Michael Yang's avatar
Michael Yang committed
348
		}
349

Michael Yang's avatar
Michael Yang committed
350
351
		if retval := C.llama_eval(llm.ctx, unsafe.SliceData(llm.embd[llm.cursor:]), C.int(numEval), C.int(llm.cursor), C.int(llm.NumThread)); retval != 0 {
			return 0, fmt.Errorf("llama_eval: %d", retval)
352
		}
Michael Yang's avatar
Michael Yang committed
353

Michael Yang's avatar
Michael Yang committed
354
		llm.cursor += numEval
Jeffrey Morgan's avatar
Jeffrey Morgan committed
355
356
	}

Michael Yang's avatar
Michael Yang committed
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
	var sampleOpts C.struct_llama_sample_options
	sampleOpts.repeat_penalty = C.float(llm.RepeatPenalty)
	sampleOpts.frequency_penalty = C.float(llm.FrequencyPenalty)
	sampleOpts.presence_penalty = C.float(llm.PresencePenalty)
	sampleOpts.temperature = C.float(llm.Temperature)
	sampleOpts.top_k = C.int(llm.TopK)
	sampleOpts.top_p = C.float(llm.TopP)
	sampleOpts.tfs_z = C.float(llm.TFSZ)
	sampleOpts.typical_p = C.float(llm.TypicalP)
	sampleOpts.mirostat = C.int(llm.Mirostat)
	sampleOpts.mirostat_tau = C.float(llm.MirostatTau)
	sampleOpts.mirostat_eta = C.float(llm.MirostatEta)
	sampleOpts.penalize_newline = C.bool(llm.PenalizeNewline)

	numVocab := C.llama_n_vocab(llm.ctx)
Michael Yang's avatar
Michael Yang committed
372
373
	logits := unsafe.Slice(C.llama_get_logits(llm.ctx), numVocab)

Michael Yang's avatar
Michael Yang committed
374
375
376
377
378
	// TODO: logit bias

	candidates := make([]C.llama_token_data, numVocab)
	for i := range logits {
		candidates[i] = C.llama_token_data{
Michael Yang's avatar
Michael Yang committed
379
380
381
			id:    C.int(i),
			logit: logits[i],
			p:     0,
Michael Yang's avatar
Michael Yang committed
382
		}
Michael Yang's avatar
Michael Yang committed
383
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
384

Michael Yang's avatar
Michael Yang committed
385
386
387
388
389
390
391
392
393
394
395
	repeatLastN := llm.RepeatLastN
	if len(llm.last) < repeatLastN {
		repeatLastN = len(llm.last)
	}

	if llm.NumCtx < repeatLastN {
		repeatLastN = llm.NumCtx
	}

	lastN := llm.last[len(llm.last)-repeatLastN:]

Michael Yang's avatar
Michael Yang committed
396
397
	token := C.llama_sample(
		llm.ctx,
Michael Yang's avatar
Michael Yang committed
398
399
400
401
402
403
404
405
406
407
		unsafe.SliceData(candidates), C.size_t(len(candidates)),
		unsafe.SliceData(lastN), C.size_t(len(lastN)),
		&sampleOpts,
	)

	llm.last = append(llm.last, token)
	llm.embd = append(llm.embd, token)

	if token == C.llama_token_eos() {
		return 0, io.EOF
Jeffrey Morgan's avatar
Jeffrey Morgan committed
408
	}
Michael Yang's avatar
Michael Yang committed
409

Michael Yang's avatar
Michael Yang committed
410
	return token, nil
Jeffrey Morgan's avatar
Jeffrey Morgan committed
411
}
412
413
414
415
416
417

func (llm *LLM) Embedding(input string) ([]float64, error) {
	if !llm.EmbeddingOnly {
		return nil, errors.New("llama: embedding not enabled")
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
418
	tokens := llm.Encode(input)
419
420
421
422
	if tokens == nil {
		return nil, errors.New("llama: tokenize embedding")
	}

423
	retval := C.llama_eval(llm.ctx, unsafe.SliceData(tokens), C.int(len(tokens)), 0, C.int(llm.NumThread))
424
425
426
427
	if retval != 0 {
		return nil, errors.New("llama: eval")
	}

428
	n := C.llama_n_embd(llm.ctx)
429
430
431
	if n <= 0 {
		return nil, errors.New("llama: no embeddings generated")
	}
432
	cEmbeddings := unsafe.Slice(C.llama_get_embeddings(llm.ctx), n)
433

434
435
436
	embeddings := make([]float64, len(cEmbeddings))
	for i, v := range cEmbeddings {
		embeddings[i] = float64(v)
437
	}
438
	return embeddings, nil
439
}