llama.go 4.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
package llm

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
10
	"sync"
11
12
13
	"time"

	"github.com/jmorganca/ollama/api"
Michael Yang's avatar
Michael Yang committed
14
	"github.com/jmorganca/ollama/format"
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
const jsonGrammar = `
root   ::= object
value  ::= object | array | string | number | ("true" | "false" | "null") ws

object ::=
  "{" ws (
            string ":" ws value
    ("," ws string ":" ws value)*
  )? "}" ws

array  ::=
  "[" ws (
            value
    ("," ws value)*
  )? "]" ws

string ::=
  "\"" (
    [^"\\] |
    "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  )* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

# Optional space: by convention, applied in this grammar after literal chars when allowed
ws ::= ([ \t\n] ws)?
`

45
46
47
48
type llamaModel struct {
	hyperparameters llamaHyperparameters
}

Michael Yang's avatar
Michael Yang committed
49
50
func (llm *llamaModel) ModelFamily() string {
	return "llama"
51
52
}

Michael Yang's avatar
Michael Yang committed
53
54
func llamaModelType(numLayer uint32) string {
	switch numLayer {
55
	case 26:
Michael Yang's avatar
Michael Yang committed
56
		return "3B"
57
	case 32:
Michael Yang's avatar
Michael Yang committed
58
		return "7B"
59
	case 40:
Michael Yang's avatar
Michael Yang committed
60
		return "13B"
61
	case 48:
Michael Yang's avatar
Michael Yang committed
62
		return "34B"
63
	case 60:
Michael Yang's avatar
Michael Yang committed
64
		return "30B"
65
	case 80:
Michael Yang's avatar
Michael Yang committed
66
67
		return "65B"
	default:
Michael Yang's avatar
Michael Yang committed
68
		return "unknown"
69
	}
Michael Yang's avatar
Michael Yang committed
70
}
71

Michael Yang's avatar
Michael Yang committed
72
73
func (llm *llamaModel) ModelType() string {
	return llamaModelType(llm.hyperparameters.NumLayer)
74
75
}

Michael Yang's avatar
Michael Yang committed
76
77
func (llm *llamaModel) FileType() string {
	return fileType(llm.hyperparameters.FileType)
78
79
}

80
81
82
83
func (llm *llamaModel) NumLayers() int64 {
	return int64(llm.hyperparameters.NumLayer)
}

84
85
86
87
88
89
90
91
92
93
94
95
96
97
type llamaHyperparameters struct {
	// NumVocab is the size of the model's vocabulary.
	NumVocab uint32

	// NumEmbd is the size of the model's embedding layer.
	NumEmbd uint32
	NumMult uint32
	NumHead uint32

	// NumLayer is the number of layers in the model.
	NumLayer uint32
	NumRot   uint32

	// FileType describes the quantization level of the model, e.g. Q4_0, Q5_K, etc.
Michael Yang's avatar
Michael Yang committed
98
	FileType uint32
99
100
101
}

type Running struct {
102
103
104
105
106
107
	Port          int
	Cmd           *exec.Cmd
	Cancel        context.CancelFunc
	exitOnce      sync.Once
	exitCh        chan error // channel to receive the exit status of the subprocess
	*StatusWriter            // captures error messages from the llama runner process
108
109
}

Patrick Devine's avatar
Patrick Devine committed
110
111
112
113
114
type ImageData struct {
	Data []byte `json:"data"`
	ID   int    `json:"id"`
}

115
116
type llama struct {
	api.Options
Patrick Devine's avatar
Patrick Devine committed
117
	ImageData []ImageData
118
119
120
	Running
}

121
var (
Jeffrey Morgan's avatar
Jeffrey Morgan committed
122
	errNvidiaSMI     = errors.New("warning: gpu support may not be enabled, check that you have installed GPU drivers: nvidia-smi command failed")
123
124
	errAvailableVRAM = errors.New("not enough VRAM available, falling back to CPU only")
)
125

126
127
// StatusWriter is a writer that captures error messages from the llama runner process
type StatusWriter struct {
128
129
	ErrCh      chan error
	LastErrMsg string
130
131
132
133
134
135
136
137
138
}

func NewStatusWriter() *StatusWriter {
	return &StatusWriter{
		ErrCh: make(chan error, 1),
	}
}

func (w *StatusWriter) Write(b []byte) (int, error) {
139
	var errMsg string
140
	if _, after, ok := bytes.Cut(b, []byte("error:")); ok {
141
142
143
		errMsg = string(bytes.TrimSpace(after))
	} else if _, after, ok := bytes.Cut(b, []byte("CUDA error")); ok {
		errMsg = string(bytes.TrimSpace(after))
144
	}
145
146
147
148
149
150

	if errMsg != "" {
		w.LastErrMsg = errMsg
		w.ErrCh <- fmt.Errorf("llama runner: %s", errMsg)
	}

151
152
153
	return os.Stderr.Write(b)
}

Michael Yang's avatar
Michael Yang committed
154
type prediction struct {
Michael Yang's avatar
Michael Yang committed
155
156
157
158
159
	Content string `json:"content"`
	Model   string `json:"model"`
	Prompt  string `json:"prompt"`
	Stop    bool   `json:"stop"`

Michael Yang's avatar
Michael Yang committed
160
161
162
163
164
165
	Timings struct {
		PredictedN  int     `json:"predicted_n"`
		PredictedMS float64 `json:"predicted_ms"`
		PromptN     int     `json:"prompt_n"`
		PromptMS    float64 `json:"prompt_ms"`
	}
166
167
}

Michael Yang's avatar
Michael Yang committed
168
const maxBufferSize = 512 * format.KiloByte
Bruce MacDonald's avatar
Bruce MacDonald committed
169
const maxRetries = 6
170

Bruce MacDonald's avatar
Bruce MacDonald committed
171
type PredictOpts struct {
172
173
174
	Prompt string
	Format string
	Images []api.ImageData
Bruce MacDonald's avatar
Bruce MacDonald committed
175
}
176

Bruce MacDonald's avatar
Bruce MacDonald committed
177
178
179
180
181
182
183
184
type PredictResult struct {
	Content            string
	Done               bool
	PromptEvalCount    int
	PromptEvalDuration time.Duration
	EvalCount          int
	EvalDuration       time.Duration
}
185

186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
type TokenizeRequest struct {
	Content string `json:"content"`
}

type TokenizeResponse struct {
	Tokens []int `json:"tokens"`
}

type DetokenizeRequest struct {
	Tokens []int `json:"tokens"`
}

type DetokenizeResponse struct {
	Content string `json:"content"`
}

type EmbeddingRequest struct {
	Content string `json:"content"`
}

type EmbeddingResponse struct {
	Embedding []float64 `json:"embedding"`
}