Is it possible to pass logit-bias to llama.cpp?

[Madd0g]

I need a little more control over which tokens are generated. I saw discussions in llama.cpp about logit-bias, is it possible to add support for that?

It would be nice to have facilities to:

1. get logits of tokens
2. pass them to completion

Is this supported or easy to add?

Thanks

[giladgd]

@Madd0g I've added support for logit bias here; it will be available in the next beta version.

I don't want to expose the internal values and flow of the logit usage as-is like you suggested since I aim to keep this library relatively high-level.

[giladgd]

The discussion you linked talks about getting the token numbers so you can use them to set biases using a CLI flag; these are token numbers, not logits.

The API I exposed should allow you to achieve this without having to investigate the tokenizer at all.
I haven't released a new version with it yet due to a bug (that needs to be fixed before a new version is published), but I'll attach here an example so you can provide feedback before I update the version 3 beta PR with code examples.

Here is an example of how you can increase the probability of the word "hello" being generated and prevent the word "day" from being generated (using the next beta version):

import {fileURLToPath} from "url";
import path from "path";
import {getLlama, LlamaChatSession, TokenBias} from "node-llama-cpp";

const __dirname = path.dirname(fileURLToPath(import.meta.url));

const llama = await getLlama();
const model = await llama.loadModel({
    modelPath: path.join(__dirname, "models", "dolphin-2.1-mistral-7b.Q4_K_M.gguf")
});
const context = await model.createContext();
const session = new LlamaChatSession({
    contextSequence: context.getSequence()
});


const q1 = "Hi there, how are you?";
console.log("User: " + q1);

const a1 = await session.prompt(q1, {
    tokenBias: (new TokenBias(model))
        .set("Hello", 1)
        .set("hello", 1)
        .set("Day", "never")
        .set("day", "never")
        .set(model.tokenize("day"), "never") // you can also do this to set bias for specific tokens
});
console.log("AI: " + a1);

From my investigations, this method covers all the relevant tokens when using a Llama 2 model.
It's worth noting that for each bias you set using TokenBias, it tokenizes it for you both with leading whitespace and without leading space to cover both variations, so you don't have to do that manually.

To be more precise and ensure you cover all the tokens you deem relevant, you can iterate over the tokenizer metadata of a model and find the relevant tokens you'd like to set a bias for.
Here's an example of how you can validate that the above approach covers all the variations of the word "hello":

const word = "hello";

const capitalize = (str: string) => ((str[0]?.toUpperCase() ?? "") + str.slice(1));
const metadataTokens = model.fileInfo.metadata.tokenizer.ggml.tokens.reduce((items, item, index) => {
    if (item.toLowerCase().includes(word))
        return items.concat([[item, [index]]]);

    return items;
}, [] as [string, [number]][]);

console.log("Metadata tokenizer values", metadataTokens);
console.log("Model tokenized values", [
    [capitalize(word), model.tokenize(capitalize(word), false, "trimLeadingSpace")],
    [word, model.tokenize(word, false, "trimLeadingSpace")],
    [capitalize(word), model.tokenize(capitalize(word))],
    [word, model.tokenize(word)]
]);

The output from both logs should contain the same tokens.

[Madd0g]

oh I understand now, looks great, can't wait to try it

[giladgd]

@Madd0g I've released a new beta version with token bias support

[Madd0g]

great, thank you, I will experiment with it

[Madd0g]

tested both methods (passing strings and direct tokens) and it seems to work well, thank you!