Reserve 3 extra tokens for each chat completion

This commit is contained in:
Cohee 2023-09-11 17:22:31 +03:00
parent 74d627f674
commit 6c8bd06308
1 changed files with 6 additions and 2 deletions

View File

@ -576,6 +576,7 @@ function populateChatCompletion(prompts, chatCompletion, { bias, quietPrompt, ty
chatCompletion.add(collection, index);
};
chatCompletion.reserveBudget(3); // every reply is primed with <|start|>assistant<|message|>
// Character and world information
addToChatCompletion('worldInfoBefore');
addToChatCompletion('main');
@ -1761,9 +1762,12 @@ class ChatCompletion {
/**
* Reserves the tokens required by the given message from the token budget.
*
* @param {Message|MessageCollection} message - The message whose tokens to reserve.
* @param {Message|MessageCollection|number} message - The message whose tokens to reserve.
*/
reserveBudget(message) { this.decreaseTokenBudgetBy(message.getTokens()) };
reserveBudget(message) {
const tokens = typeof message === 'number' ? message : message.getTokens();
this.decreaseTokenBudgetBy(tokens);
};
/**
* Frees up the tokens used by the given message from the token budget.