Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit adbaa8b3 by archer
Options

feat: use Tiktokenizer to count tokens

parent 29c95d24
Showing with 149 additions and 32 deletions
.env.template
......@@ -3,6 +3,7 @@ AXIOS_PROXY_HOST=127.0.0.1
AXIOS_PROXY_PORT_FAST=7890
AXIOS_PROXY_PORT_NORMAL=7890
queueTask=1
parentUrl=https://hostname/api/openapi/startEvents
# email
MY_MAIL=xxx@qq.com
MAILE_CODE=xxx
......
README.md
......@@ -14,6 +14,7 @@ AXIOS_PROXY_HOST=127.0.0.1
AXIOS_PROXY_PORT_FAST=7890
AXIOS_PROXY_PORT_NORMAL=7890
queueTask=1
parentUrl=https://hostname/api/openapi/startEvents
# email，参考 nodeMail 获取参数
MY_MAIL=xxx@qq.com
MAILE_CODE=xxx
......@@ -196,7 +197,7 @@ services:
- TOKEN_KEY=xxxx
# 是否开启队列任务。 1-开启，0-关闭（请求parentUrl去执行任务,单机时直接填1）
- queueTask=1
- parentUrl=https://fastgpt.run/api/openapi/startEvents
- parentUrl=https://hostname/api/openapi/startEvents
# db
- MONGODB_URI=mongodb://username:passsword@0.0.0.0:27017/?authSource=admin
- MONGODB_NAME=xxx
......
next.config.js
......@@ -4,7 +4,12 @@ const nextConfig = {
output: 'standalone',
reactStrictMode: true,
compress: true,
webpack(config) {
config.experiments = {
asyncWebAssembly: true,
layers: true
};
config.module.rules = config.module.rules.concat([
{
test: /\.svg$/i,
......
package.json
......@@ -17,6 +17,7 @@
"@chakra-ui/icons": "^2.0.17",
"@chakra-ui/react": "^2.5.1",
"@chakra-ui/system": "^2.5.5",
"@dqbd/tiktoken": "^1.0.6",
"@emotion/react": "^11.10.6",
"@emotion/styled": "^11.10.6",
"@next/font": "13.1.6",
......@@ -29,6 +30,7 @@
"formidable": "^2.1.1",
"framer-motion": "^9.0.6",
"gpt-token-utils": "^1.2.0",
"graphemer": "^1.4.0",
"hyperdown": "^2.4.29",
"immer": "^9.0.19",
"jsonwebtoken": "^9.0.0",
......
pnpm-lock.yaml
......@@ -7,6 +7,7 @@ specifiers:
'@chakra-ui/icons': ^2.0.17
'@chakra-ui/react': ^2.5.1
'@chakra-ui/system': ^2.5.5
'@dqbd/tiktoken': ^1.0.6
'@emotion/react': ^11.10.6
'@emotion/styled': ^11.10.6
'@next/font': 13.1.6
......@@ -33,6 +34,7 @@ specifiers:
formidable: ^2.1.1
framer-motion: ^9.0.6
gpt-token-utils: ^1.2.0
graphemer: ^1.4.0
husky: ^8.0.3
hyperdown: ^2.4.29
immer: ^9.0.19
......@@ -72,6 +74,7 @@ dependencies:
'@chakra-ui/icons': registry.npmmirror.com/@chakra-ui/icons/2.0.17_lze4h7kxffpjhokvtqbtrlfkmq
'@chakra-ui/react': registry.npmmirror.com/@chakra-ui/react/2.5.1_e6pzu3hsaqmql4fl7jx73ckiym
'@chakra-ui/system': registry.npmmirror.com/@chakra-ui/system/2.5.5_xqp3pgpqjlfxxa3zxu4zoc4fba
'@dqbd/tiktoken': registry.npmmirror.com/@dqbd/tiktoken/1.0.6
'@emotion/react': registry.npmmirror.com/@emotion/react/11.10.6_pmekkgnqduwlme35zpnqhenc34
'@emotion/styled': registry.npmmirror.com/@emotion/styled/11.10.6_oouaibmszuch5k64ms7uxp2aia
'@next/font': registry.npmmirror.com/@next/font/13.1.6
......@@ -84,6 +87,7 @@ dependencies:
formidable: registry.npmmirror.com/formidable/2.1.1
framer-motion: registry.npmmirror.com/framer-motion/9.0.6_biqbaboplfbrettd7655fr4n2y
gpt-token-utils: registry.npmmirror.com/gpt-token-utils/1.2.0
graphemer: registry.npmmirror.com/graphemer/1.4.0
hyperdown: registry.npmmirror.com/hyperdown/2.4.29
immer: registry.npmmirror.com/immer/9.0.19
jsonwebtoken: registry.npmmirror.com/jsonwebtoken/9.0.0
......@@ -4176,6 +4180,12 @@ packages:
react: registry.npmmirror.com/react/18.2.0
dev: false
registry.npmmirror.com/@dqbd/tiktoken/1.0.6:
resolution: {integrity: sha512-umSdeZTy/SbPPKVuZKV/XKyFPmXSN145CcM3iHjBbmhlohBJg7vaDp4cPCW+xNlWL6L2U1sp7T2BD+di2sUKdA==, registry: https://registry.npm.taobao.org/, tarball: https://registry.npmmirror.com/@dqbd/tiktoken/-/tiktoken-1.0.6.tgz}
name: '@dqbd/tiktoken'
version: 1.0.6
dev: false
registry.npmmirror.com/@emotion/babel-plugin/11.10.6:
resolution: {integrity: sha512-p2dAqtVrkhSa7xz1u/m9eHYdLi+en8NowrmXeF/dKtJpU8lCWli8RUAati7NcSl0afsBott48pdnANuD0wh9QQ==, registry: https://registry.npm.taobao.org/, tarball: https://registry.npmmirror.com/@emotion/babel-plugin/-/babel-plugin-11.10.6.tgz}
name: '@emotion/babel-plugin'
......@@ -7675,6 +7685,12 @@ packages:
version: 1.0.4
dev: true
registry.npmmirror.com/graphemer/1.4.0:
resolution: {integrity: sha512-EtKwoO6kxCL9WO5xipiHTZlSzBm7WLT627TqC/uVRd0HKmq8NXyebnNYxDoBi7wt8eTWrUrKXCOVaFq9x1kgag==, registry: https://registry.npm.taobao.org/, tarball: https://registry.npmmirror.com/graphemer/-/graphemer-1.4.0.tgz}
name: graphemer
version: 1.4.0
dev: false
registry.npmmirror.com/has-bigints/1.0.2:
resolution: {integrity: sha512-tSvCKtBr9lkF0Ex0aQiP9N+OpV4zi2r/Nee5VkRDbaqv35RLYMzbwQfFSZZH0kR+Rd6302UJZ2p/bJCEoR3VoQ==, registry: https://registry.npm.taobao.org/, tarball: https://registry.npmmirror.com/has-bigints/-/has-bigints-1.0.2.tgz}
name: has-bigints
......
src/pages/model/detail/components/SelectFileModal.tsx
......@@ -14,7 +14,6 @@ import {
} from '@chakra-ui/react';
import { useToast } from '@/hooks/useToast';
import { useSelectFile } from '@/hooks/useSelectFile';
import { encode } from 'gpt-token-utils';
import { useConfirm } from '@/hooks/useConfirm';
import { readTxtContent, readPdfContent, readDocContent } from '@/utils/file';
import { useMutation } from '@tanstack/react-query';
......@@ -22,6 +21,7 @@ import { postModelDataSplitData } from '@/api/model';
import { formatPrice } from '@/utils/user';
import Radio from '@/components/Radio';
import { splitText } from '@/utils/file';
import { countChatTokens } from '@/utils/tools';
const fileExtension = '.txt,.doc,.docx,.pdf,.md';
......@@ -29,11 +29,11 @@ const modeMap = {
qa: {
maxLen: 2800,
slideLen: 800,
price: 3,
price: 4,
isPrompt: true
},
subsection: {
maxLen: 1000,
maxLen: 800,
slideLen: 300,
price: 0.4,
isPrompt: false
......@@ -55,19 +55,19 @@ const SelectFileModal = ({
const { File, onOpen } = useSelectFile({ fileType: fileExtension, multiple: true });
const [mode, setMode] = useState<'qa' | 'subsection'>('qa');
const [fileTextArr, setFileTextArr] = useState<string[]>(['']);
const [splitRes, setSplitRes] = useState<{ tokens: number; chunks: string[] }>({
tokens: 0,
chunks: []
});
const { openConfirm, ConfirmChild } = useConfirm({
content: '确认导入该文件，需要一定时间进行拆解，该任务无法终止！如果余额不足，任务讲被终止。'
content: `确认导入该文件，需要一定时间进行拆解，该任务无法终止！如果余额不足，未完成的任务会被直接清除。一共 ${
splitRes.chunks.length
} 组，大约 ${splitRes.tokens} 个tokens, 约 ${formatPrice(
splitRes.tokens * modeMap[mode].price
)} 元`
});
const fileText = useMemo(() => {
const chunks = fileTextArr.map((item) =>
splitText({
text: item,
...modeMap[mode]
})
);
return chunks.join('');
}, [fileTextArr, mode]);
const fileText = useMemo(() => fileTextArr.join(''), [fileTextArr]);
const onSelectFile = useCallback(
async (e: File[]) => {
......@@ -106,18 +106,11 @@ const SelectFileModal = ({
const { mutate, isLoading } = useMutation({
mutationFn: async () => {
if (!fileText) return;
const chunks = fileTextArr
.map((item) =>
splitText({
text: item,
...modeMap[mode]
})
)
.flat();
if (splitRes.chunks.length === 0) return;
await postModelDataSplitData({
modelId,
chunks,
chunks: splitRes.chunks,
prompt: `下面是"${prompt || '一段长文本'}"`,
mode
});
......@@ -136,6 +129,28 @@ const SelectFileModal = ({
}
});
const onclickImport = useCallback(() => {
const chunks = fileTextArr
.map((item) =>
splitText({
text: item,
...modeMap[mode]
})
)
.flat();
// count tokens
const tokens = chunks.map((item) =>
countChatTokens({ messages: [{ role: 'system', content: item }] })
);
setSplitRes({
tokens: tokens.reduce((sum, item) => sum + item, 0),
chunks
});
openConfirm(mutate)();
}, [fileTextArr, mode, mutate, openConfirm]);
return (
<Modal isOpen={true} onClose={onClose} isCentered>
<ModalOverlay />
......@@ -152,10 +167,9 @@ const SelectFileModal = ({
justifyContent={'center'}
fontSize={'sm'}
>
<Box mt={2} px={4} maxW={['100%']} textAlign={'justify'} color={'blackAlpha.600'}>
<Box mt={2} px={5} maxW={['100%', '70%']} textAlign={'justify'} color={'blackAlpha.600'}>
支持 {fileExtension} 文件。模型会自动对文本进行 QA 拆分，需要较长训练时间，拆分需要消耗
tokens，账号余额不足时，未拆分的数据会被删除。当前一共 {encode(fileText).length}{' '}
个tokens，大约 {formatPrice(encode(fileText).length * modeMap[mode].price)}元
tokens，账号余额不足时，未拆分的数据会被删除。
</Box>
{/* 拆分模式 */}
<Flex w={'100%'} px={5} alignItems={'center'} mt={4}>
......@@ -217,7 +231,7 @@ const SelectFileModal = ({
<Button variant={'outline'} colorScheme={'gray'} mr={3} onClick={onClose}>
取消
</Button>
<Button isLoading={isLoading} isDisabled={fileText === ''} onClick={openConfirm(mutate)}>
<Button isLoading={isLoading} isDisabled={fileText === ''} onClick={onclickImport}>
确认导入
</Button>
</Flex>
......
src/utils/file.ts
import mammoth from 'mammoth';
import Papa from 'papaparse';
import { encode } from 'gpt-token-utils';
import { countChatTokens } from './tools';
/**
* 读取 txt 文件内容
......@@ -164,7 +164,7 @@ export const splitText = ({
const chunks: { sum: number; arr: string[] }[] = [{ sum: 0, arr: [] }];
for (let i = 0; i < textArr.length; i++) {
const tokenLen = encode(textArr[i]).length;
const tokenLen = countChatTokens({ messages: [{ role: 'system', content: textArr[i] }] });
chunks[chunks.length - 1].sum += tokenLen;
chunks[chunks.length - 1].arr.push(textArr[i]);
......@@ -174,7 +174,7 @@ export const splitText = ({
const chunk: { sum: number; arr: string[] } = { sum: 0, arr: [] };
for (let j = chunks[chunks.length - 1].arr.length - 1; j >= 0; j--) {
const chunkText = chunks[chunks.length - 1].arr[j];
const tokenLen = encode(chunkText).length;
const tokenLen = countChatTokens({ messages: [{ role: 'system', content: chunkText }] });
chunk.sum += tokenLen;
chunk.arr.unshift(chunkText);
......@@ -185,7 +185,6 @@ export const splitText = ({
chunks.push(chunk);
}
}
const result = chunks.map((item) => item.arr.join(''));
return result;
};
src/utils/tools.ts
import crypto from 'crypto';
import { useToast } from '@/hooks/useToast';
import { encoding_for_model, type Tiktoken } from '@dqbd/tiktoken';
import Graphemer from 'graphemer';
const textDecoder = new TextDecoder();
const graphemer = new Graphemer();
const encMap = {
'gpt-3.5-turbo': encoding_for_model('gpt-3.5-turbo', {
'<|im_start|>': 100264,
'<|im_end|>': 100265,
'<|im_sep|>': 100266
}),
'gpt-4': encoding_for_model('gpt-4', {
'<|im_start|>': 100264,
'<|im_end|>': 100265,
'<|im_sep|>': 100266
}),
'gpt-4-32k': encoding_for_model('gpt-4-32k', {
'<|im_start|>': 100264,
'<|im_end|>': 100265,
'<|im_sep|>': 100266
})
};
/**
* copy text data
......@@ -51,3 +73,60 @@ export const Obj2Query = (obj: Record<string, string | number>) => {
}
return queryParams.toString();
};
/* 格式化 chat 聊天内容 */
function getChatGPTEncodingText(
messages: { role: 'system' | 'user' | 'assistant'; content: string; name?: string }[],
model: 'gpt-3.5-turbo' | 'gpt-4' | 'gpt-4-32k'
) {
const isGpt3 = model === 'gpt-3.5-turbo';
const msgSep = isGpt3 ? '\n' : '';
const roleSep = isGpt3 ? '\n' : '<|im_sep|>';
return [
messages
.map(({ name = '', role, content }) => {
return `<|im_start|>${name || role}${roleSep}${content}<|im_end|>`;
})
.join(msgSep),
`<|im_start|>assistant${roleSep}`
].join(msgSep);
}
function text2TokensLen(encoder: Tiktoken, inputText: string) {
const encoding = encoder.encode(inputText, 'all');
const segments: { text: string; tokens: { id: number; idx: number }[] }[] = [];
let byteAcc: number[] = [];
let tokenAcc: { id: number; idx: number }[] = [];
let inputGraphemes = graphemer.splitGraphemes(inputText);
for (let idx = 0; idx < encoding.length; idx++) {
const token = encoding[idx]!;
byteAcc.push(...encoder.decode_single_token_bytes(token));
tokenAcc.push({ id: token, idx });
const segmentText = textDecoder.decode(new Uint8Array(byteAcc));
const graphemes = graphemer.splitGraphemes(segmentText);
if (graphemes.every((item, idx) => inputGraphemes[idx] === item)) {
segments.push({ text: segmentText, tokens: tokenAcc });
byteAcc = [];
tokenAcc = [];
inputGraphemes = inputGraphemes.slice(graphemes.length);
}
}
return segments.reduce((memo, i) => memo + i.tokens.length, 0) ?? 0;
}
export const countChatTokens = ({
model = 'gpt-3.5-turbo',
messages
}: {
model?: 'gpt-4' | 'gpt-4-32k' | 'gpt-3.5-turbo';
messages: { role: 'system' | 'user' | 'assistant'; content: string }[];
}) => {
const text = getChatGPTEncodingText(messages, model);
return text2TokensLen(encMap[model], text);
};
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment