Skip to content

Commit ede69c1

Browse files
giladgdgiladgd
authored
feat: completion and infill (#164)
* feat: add `LlamaCompletion` that provides the ability to complete or infill text * feat: support configuring more options for `getLlama` when using `"lastBuild"` * fix: various bug fixes
1 parent 47b476f commit ede69c1

26 files changed

+2054
-124
lines changed

‎docs/guide/cli/cli.data.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ import {CommandModule} from "yargs";
22
import {getCommandHtmlDoc} from "../../../.vitepress/utils/getCommandHtmlDoc.js";
33
import {BuildCommand} from "../../../src/cli/commands/BuildCommand.js";
44
import {ChatCommand} from "../../../src/cli/commands/ChatCommand.js";
5+
import {CompleteCommand} from "../../../src/cli/commands/CompleteCommand.js";
6+
import {InfillCommand} from "../../../src/cli/commands/InfillCommand.js";
57
import {DownloadCommand} from "../../../src/cli/commands/DownloadCommand.js";
68
import {ClearCommand} from "../../../src/cli/commands/ClearCommand.js";
79
import {htmlEscape} from "../../../.vitepress/utils/htmlEscape.js";
@@ -17,12 +19,16 @@ export default {
1719
return {
1820
index: buildIndexTable([
1921
["chat", ChatCommand],
22+
["complete", CompleteCommand],
23+
["infill", InfillCommand],
2024
["download", DownloadCommand],
2125
["build", BuildCommand],
2226
["clear", ClearCommand]
2327
]),
2428

2529
chat: await getCommandHtmlDoc(ChatCommand),
30+
complete: await getCommandHtmlDoc(CompleteCommand),
31+
infill: await getCommandHtmlDoc(InfillCommand),
2632
download: await getCommandHtmlDoc(DownloadCommand),
2733
build: await getCommandHtmlDoc(BuildCommand),
2834
clear: await getCommandHtmlDoc(ClearCommand)

‎docs/guide/cli/complete.md‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
---
2+
outline: deep
3+
---
4+
# `complete` command
5+
6+
<script setup lang="ts">
7+
import {data as docs} from "./cli.data.js";
8+
const commandDoc = docs.complete;
9+
</script>
10+
11+
{{commandDoc.description}}
12+
13+
## Usage
14+
```shell-vue
15+
{{commandDoc.usage}}
16+
```
17+
<div v-html="commandDoc.options"></div>

‎docs/guide/cli/infill.md‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
---
2+
outline: deep
3+
---
4+
# `infill` command
5+
6+
<script setup lang="ts">
7+
import {data as docs} from "./cli.data.js";
8+
const commandDoc = docs.infill;
9+
</script>
10+
11+
{{commandDoc.description}}
12+
13+
## Usage
14+
```shell-vue
15+
{{commandDoc.usage}}
16+
```
17+
<div v-html="commandDoc.options"></div>

‎llama/addon.cpp‎

Lines changed: 63 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,26 @@ Napi::Value getGpuVramInfo(const Napi::CallbackInfo& info) {
8787
return result;
8888
}
8989

90+
static Napi::Value getNapiToken(const Napi::CallbackInfo& info, llama_model* model, llama_token token) {
91+
auto tokenType = llama_token_get_type(model, token);
92+
93+
if (tokenType == LLAMA_TOKEN_TYPE_UNDEFINED || tokenType == LLAMA_TOKEN_TYPE_UNKNOWN) {
94+
return Napi::Number::From(info.Env(), -1);
95+
}
96+
97+
return Napi::Number::From(info.Env(), token);
98+
}
99+
100+
static Napi::Value getNapiControlToken(const Napi::CallbackInfo& info, llama_model* model, llama_token token) {
101+
auto tokenType = llama_token_get_type(model, token);
102+
103+
if (tokenType != LLAMA_TOKEN_TYPE_CONTROL) {
104+
return Napi::Number::From(info.Env(), -1);
105+
}
106+
107+
return Napi::Number::From(info.Env(), token);
108+
}
109+
90110
class AddonModel : public Napi::ObjectWrap<AddonModel> {
91111
public:
92112
llama_model_params model_params;
@@ -119,7 +139,6 @@ class AddonModel : public Napi::ObjectWrap<AddonModel> {
119139
}
120140
}
121141

122-
llama_backend_init(false);
123142
model = llama_load_model_from_file(modelPath.c_str(), model_params);
124143

125144
if (model == NULL) {
@@ -203,6 +222,15 @@ class AddonModel : public Napi::ObjectWrap<AddonModel> {
203222
return Napi::Number::From(info.Env(), llama_n_ctx_train(model));
204223
}
205224

225+
Napi::Value GetEmbeddingVectorSize(const Napi::CallbackInfo& info) {
226+
if (disposed) {
227+
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
228+
return info.Env().Undefined();
229+
}
230+
231+
return Napi::Number::From(info.Env(), llama_n_embd(model));
232+
}
233+
206234
Napi::Value GetTotalSize(const Napi::CallbackInfo& info) {
207235
if (disposed) {
208236
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
@@ -239,55 +267,55 @@ class AddonModel : public Napi::ObjectWrap<AddonModel> {
239267
return info.Env().Undefined();
240268
}
241269

242-
return Napi::Number::From(info.Env(), llama_token_bos(model));
270+
return getNapiControlToken(info, model, llama_token_bos(model));
243271
}
244272
Napi::Value TokenEos(const Napi::CallbackInfo& info) {
245273
if (disposed) {
246274
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
247275
return info.Env().Undefined();
248276
}
249277

250-
return Napi::Number::From(info.Env(), llama_token_eos(model));
278+
return getNapiControlToken(info, model, llama_token_eos(model));
251279
}
252280
Napi::Value TokenNl(const Napi::CallbackInfo& info) {
253281
if (disposed) {
254282
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
255283
return info.Env().Undefined();
256284
}
257285

258-
return Napi::Number::From(info.Env(), llama_token_nl(model));
286+
return getNapiToken(info, model, llama_token_nl(model));
259287
}
260288
Napi::Value PrefixToken(const Napi::CallbackInfo& info) {
261289
if (disposed) {
262290
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
263291
return info.Env().Undefined();
264292
}
265293

266-
return Napi::Number::From(info.Env(), llama_token_prefix(model));
294+
return getNapiControlToken(info, model, llama_token_prefix(model));
267295
}
268296
Napi::Value MiddleToken(const Napi::CallbackInfo& info) {
269297
if (disposed) {
270298
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
271299
return info.Env().Undefined();
272300
}
273301

274-
return Napi::Number::From(info.Env(), llama_token_middle(model));
302+
return getNapiControlToken(info, model, llama_token_middle(model));
275303
}
276304
Napi::Value SuffixToken(const Napi::CallbackInfo& info) {
277305
if (disposed) {
278306
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
279307
return info.Env().Undefined();
280308
}
281309

282-
return Napi::Number::From(info.Env(), llama_token_suffix(model));
310+
return getNapiControlToken(info, model, llama_token_suffix(model));
283311
}
284312
Napi::Value EotToken(const Napi::CallbackInfo& info) {
285313
if (disposed) {
286314
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
287315
return info.Env().Undefined();
288316
}
289317

290-
return Napi::Number::From(info.Env(), llama_token_eot(model));
318+
return getNapiControlToken(info, model, llama_token_eot(model));
291319
}
292320
Napi::Value GetTokenString(const Napi::CallbackInfo& info) {
293321
if (disposed) {
@@ -308,6 +336,29 @@ class AddonModel : public Napi::ObjectWrap<AddonModel> {
308336
return Napi::String::New(info.Env(), ss.str());
309337
}
310338

339+
Napi::Value GetTokenType(const Napi::CallbackInfo& info) {
340+
if (disposed) {
341+
Napi::Error::New(info.Env(), "Context is disposed").ThrowAsJavaScriptException();
342+
return info.Env().Undefined();
343+
}
344+
345+
if (info[0].IsNumber() == false) {
346+
return Napi::Number::From(info.Env(), int32_t(LLAMA_TOKEN_TYPE_UNDEFINED));
347+
}
348+
349+
int token = info[0].As<Napi::Number>().Int32Value();
350+
auto tokenType = llama_token_get_type(model, token);
351+
352+
return Napi::Number::From(info.Env(), int32_t(tokenType));
353+
}
354+
Napi::Value ShouldPrependBosToken(const Napi::CallbackInfo& info) {
355+
const int addBos = llama_add_bos_token(model);
356+
357+
bool shouldPrependBos = addBos != -1 ? bool(addBos) : (llama_vocab_type(model) == LLAMA_VOCAB_TYPE_SPM);
358+
359+
return Napi::Boolean::New(info.Env(), shouldPrependBos);
360+
}
361+
311362
static void init(Napi::Object exports) {
312363
exports.Set(
313364
"AddonModel",
@@ -318,6 +369,7 @@ class AddonModel : public Napi::ObjectWrap<AddonModel> {
318369
InstanceMethod("tokenize", &AddonModel::Tokenize),
319370
InstanceMethod("detokenize", &AddonModel::Detokenize),
320371
InstanceMethod("getTrainContextSize", &AddonModel::GetTrainContextSize),
372+
InstanceMethod("getEmbeddingVectorSize", &AddonModel::GetEmbeddingVectorSize),
321373
InstanceMethod("getTotalSize", &AddonModel::GetTotalSize),
322374
InstanceMethod("getTotalParameters", &AddonModel::GetTotalParameters),
323375
InstanceMethod("getModelDescription", &AddonModel::GetModelDescription),
@@ -329,6 +381,8 @@ class AddonModel : public Napi::ObjectWrap<AddonModel> {
329381
InstanceMethod("suffixToken", &AddonModel::SuffixToken),
330382
InstanceMethod("eotToken", &AddonModel::EotToken),
331383
InstanceMethod("getTokenString", &AddonModel::GetTokenString),
384+
InstanceMethod("getTokenType", &AddonModel::GetTokenType),
385+
InstanceMethod("shouldPrependBosToken", &AddonModel::ShouldPrependBosToken),
332386
InstanceMethod("dispose", &AddonModel::Dispose),
333387
}
334388
)
@@ -993,7 +1047,7 @@ Napi::Value setLoggerLogLevel(const Napi::CallbackInfo& info) {
9931047
}
9941048

9951049
Napi::Object registerCallback(Napi::Env env, Napi::Object exports) {
996-
llama_backend_init(false);
1050+
llama_backend_init();
9971051
exports.DefineProperties({
9981052
Napi::PropertyDescriptor::Function("systemInfo", systemInfo),
9991053
Napi::PropertyDescriptor::Function("setLogger", setLogger),

‎package-lock.json‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@
154154
"cross-spawn": "^7.0.3",
155155
"env-var": "^7.3.1",
156156
"fs-extra": "^11.2.0",
157-
"lifecycle-utils": "^1.2.1",
157+
"lifecycle-utils": "^1.2.2",
158158
"log-symbols": "^5.1.0",
159159
"node-addon-api": "^7.0.0",
160160
"octokit": "^3.1.0",

‎src/AbortError.ts‎

Lines changed: 0 additions & 6 deletions
This file was deleted.

‎src/bindings/AddonTypes.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ export type AddonModel = {
4242
tokenize(text: string, specialTokens: boolean): Uint32Array,
4343
detokenize(tokens: Uint32Array): string,
4444
getTrainContextSize(): number,
45+
getEmbeddingVectorSize(): number,
4546
getTotalSize(): number,
4647
getTotalParameters(): number,
4748
getModelDescription(): ModelTypeDescription,
@@ -52,7 +53,9 @@ export type AddonModel = {
5253
middleToken(): Token,
5354
suffixToken(): Token,
5455
eotToken(): Token,
55-
getTokenString(token: number): string
56+
getTokenString(token: number): string,
57+
getTokenType(token: Token): number,
58+
shouldPrependBosToken(): boolean
5659
};
5760

5861
export type AddonContext = {

‎src/bindings/getLlama.ts‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,27 @@ export type LastBuildOptions = {
106106
/**
107107
* Set a custom logger for llama.cpp logs.
108108
*/
109-
logger?: (level: LlamaLogLevel, message: string) => void
109+
logger?: (level: LlamaLogLevel, message: string) => void,
110+
111+
/**
112+
* If a local build is not found, use prebuilt binaries.
113+
* Enabled by default.
114+
*/
115+
usePrebuiltBinaries?: boolean,
116+
117+
/**
118+
* If a local build is not found, and prebuilt binaries are not found, when building from source,
119+
* print binary compilation progress logs.
120+
* Enabled by default.
121+
*/
122+
progressLogs?: boolean,
123+
124+
/**
125+
* If a local build is not found, and prebuilt binaries are not found, don't download llama.cpp source if it's not found.
126+
* When set to `true`, and llama.cpp source is needed but is not found, a `NoBinaryFoundError` error will be thrown.
127+
* Disabled by default.
128+
*/
129+
skipDownload?: boolean
110130
};
111131

112132
export const getLlamaFunctionName = "getLlama";
@@ -124,7 +144,10 @@ export async function getLlama(options?: LlamaOptions | "lastBuild", lastBuildOp
124144
const lastBuildInfo = await getLastBuildInfo();
125145
const getLlamaOptions: LlamaOptions = {
126146
logLevel: lastBuildOptions?.logLevel ?? defaultLlamaCppDebugLogs,
127-
logger: lastBuildOptions?.logger ?? Llama.defaultConsoleLogger
147+
logger: lastBuildOptions?.logger ?? Llama.defaultConsoleLogger,
148+
usePrebuiltBinaries: lastBuildOptions?.usePrebuiltBinaries ?? true,
149+
progressLogs: lastBuildOptions?.progressLogs ?? true,
150+
skipDownload: lastBuildOptions?.skipDownload ?? defaultSkipDownload
128151
};
129152

130153
if (lastBuildInfo == null)

‎src/cli/cli.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import {BuildCommand} from "./commands/BuildCommand.js";
1313
import {OnPostInstallCommand} from "./commands/OnPostInstallCommand.js";
1414
import {ClearCommand} from "./commands/ClearCommand.js";
1515
import {ChatCommand} from "./commands/ChatCommand.js";
16+
import {CompleteCommand} from "./commands/CompleteCommand.js";
17+
import {InfillCommand} from "./commands/InfillCommand.js";
1618
import {DebugCommand} from "./commands/DebugCommand.js";
1719

1820
const __dirname = path.dirname(fileURLToPath(import.meta.url));
@@ -30,6 +32,8 @@ yarg
3032
.command(BuildCommand)
3133
.command(ClearCommand)
3234
.command(ChatCommand)
35+
.command(CompleteCommand)
36+
.command(InfillCommand)
3337
.command(OnPostInstallCommand)
3438
.command(DebugCommand)
3539
.recommendCommands()

0 commit comments

Comments
 (0)