KeyArgo
/
llama-gpt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

include prompt in token count (#104)

This commit is contained in:
Alan P 2023-03-23 17:51:51 -04:00 committed by GitHub
parent 7ce2d5ec2c
commit 1a4b4401ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
pages/api/chat.ts
View File

@ -7,7 +7,7 @@ import { init, Tiktoken } from "@dqbd/tiktoken/lite/init";
import wasm from "../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module"; import wasm from "../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module";
export const config = { export const config = {
runtime: "edge" runtime: "edge",
}; };
const handler = async (req: Request): Promise<Response> => { const handler = async (req: Request): Promise<Response> => {
@ -15,10 +15,22 @@ const handler = async (req: Request): Promise<Response> => {
const { model, messages, key, prompt } = (await req.json()) as ChatBody; const { model, messages, key, prompt } = (await req.json()) as ChatBody;
await init((imports) => WebAssembly.instantiate(wasm, imports)); await init((imports) => WebAssembly.instantiate(wasm, imports));
const encoding = new Tiktoken(tiktokenModel.bpe_ranks, tiktokenModel.special_tokens, tiktokenModel.pat_str); const encoding = new Tiktoken(
tiktokenModel.bpe_ranks,
tiktokenModel.special_tokens,
tiktokenModel.pat_str
);
const tokenLimit = model.id === OpenAIModelID.GPT_4 ? 6000 : 3000; const tokenLimit = model.id === OpenAIModelID.GPT_4 ? 6000 : 3000;
let tokenCount = 0;
let promptToSend = prompt;
if (!promptToSend) {
promptToSend = DEFAULT_SYSTEM_PROMPT;
}
const prompt_tokens = encoding.encode(promptToSend);
let tokenCount = prompt_tokens.length;
let messagesToSend: Message[] = []; let messagesToSend: Message[] = [];
for (let i = messages.length - 1; i >= 0; i--) { for (let i = messages.length - 1; i >= 0; i--) {
@ -34,11 +46,6 @@ const handler = async (req: Request): Promise<Response> => {
encoding.free(); encoding.free();
let promptToSend = prompt;
if (!promptToSend) {
promptToSend = DEFAULT_SYSTEM_PROMPT;
}
const stream = await OpenAIStream(model, promptToSend, key, messagesToSend); const stream = await OpenAIStream(model, promptToSend, key, messagesToSend);
return new Response(stream); return new Response(stream);