mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-15 00:56:26 +08:00
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:
@@ -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');
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user