feat(server): use native tokenizer impl (#6960)

### Benchmark

`yarn workspace @affine/server-native bench`

```
┌─────────┬────────────┬─────────┬────────────────────┬──────────┬─────────┐
│ (index) │ Task Name  │ ops/sec │ Average Time (ns)  │ Margin   │ Samples │
├─────────┼────────────┼─────────┼────────────────────┼──────────┼─────────┤
│ 0       │ 'tiktoken' │ '5'     │ 176932518.76000002 │ '±4.71%' │ 100     │
│ 1       │ 'native'   │ '16'    │ 61041597.51000003  │ '±0.60%' │ 100     │
└─────────┴────────────┴─────────┴────────────────────┴──────────┴─────────┘
```
This commit is contained in:
Brooooooklyn
2024-05-16 07:55:10 +00:00
parent 46140039d9
commit c7ddd679fd
17 changed files with 206 additions and 39 deletions
@@ -1,7 +1,7 @@
import { type Tokenizer } from '@affine/server-native';
import { Injectable, Logger } from '@nestjs/common';
import { AiPrompt, PrismaClient } from '@prisma/client';
import Mustache from 'mustache';
import { Tiktoken } from 'tiktoken';
import {
getTokenEncoder,
@@ -27,7 +27,7 @@ function extractMustacheParams(template: string) {
export class ChatPrompt {
private readonly logger = new Logger(ChatPrompt.name);
public readonly encoder?: Tiktoken;
public readonly encoder: Tokenizer | null;
private readonly promptTokenSize: number;
private readonly templateParamKeys: string[] = [];
private readonly templateParams: PromptParams = {};
@@ -53,8 +53,7 @@ export class ChatPrompt {
) {
this.encoder = getTokenEncoder(model);
this.promptTokenSize =
this.encoder?.encode_ordinary(messages.map(m => m.content).join('') || '')
.length || 0;
this.encoder?.count(messages.map(m => m.content).join('') || '') || 0;
this.templateParamKeys = extractMustacheParams(
messages.map(m => m.content).join('')
);
@@ -86,7 +85,7 @@ export class ChatPrompt {
}
encode(message: string) {
return this.encoder?.encode_ordinary(message).length || 0;
return this.encoder?.count(message) || 0;
}
private checkParams(params: PromptParams, sessionId?: string) {
@@ -129,10 +128,6 @@ export class ChatPrompt {
content: Mustache.render(content, params),
}));
}
free() {
this.encoder?.free();
}
}
@Injectable()
@@ -164,7 +164,6 @@ export class ChatSession implements AsyncDisposable {
}
async [Symbol.asyncDispose]() {
this.state.prompt.free();
await this.save?.();
}
}
@@ -323,7 +322,7 @@ export class ChatSessionService {
): number {
const encoder = getTokenEncoder(model);
return messages
.map(m => encoder?.encode_ordinary(m.content).length || 0)
.map(m => encoder?.count(m.content) ?? 0)
.reduce((total, length) => total + length, 0);
}
@@ -1,13 +1,9 @@
import { type Tokenizer } from '@affine/server-native';
import { AiPromptRole } from '@prisma/client';
import type { ClientOptions as OpenAIClientOptions } from 'openai';
import {
encoding_for_model,
get_encoding,
Tiktoken,
TiktokenModel,
} from 'tiktoken';
import { z } from 'zod';
import { fromModelName } from '../../native';
import type { ChatPrompt } from './prompt';
import type { FalConfig } from './providers/fal';
@@ -37,17 +33,17 @@ export enum AvailableModels {
export type AvailableModel = keyof typeof AvailableModels;
export function getTokenEncoder(model?: string | null): Tiktoken | undefined {
if (!model) return undefined;
export function getTokenEncoder(model?: string | null): Tokenizer | null {
if (!model) return null;
const modelStr = AvailableModels[model as AvailableModel];
if (!modelStr) return undefined;
if (!modelStr) return null;
if (modelStr.startsWith('gpt')) {
return encoding_for_model(modelStr as TiktokenModel);
return fromModelName(modelStr);
} else if (modelStr.startsWith('dall')) {
// dalle don't need to calc the token
return undefined;
return null;
} else {
return get_encoding('cl100k_base');
return fromModelName('gpt-4-turbo-preview');
}
}