Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit 9f07741e by YeYuheng Committed by GitHub
Options

fix: dataset code block splitting (#6998)

parent 3b052779
Showing with 67 additions and 11 deletions
packages/global/common/string/textSplitter.ts
...@@ -63,7 +63,7 @@ const strIsMdTable = (str: string) => { ...@@ -63,7 +63,7 @@ const strIsMdTable = (str: string) => {
return true; return true;
}; };
const markdownTableSplit = (props: SplitProps): SplitResponse => { const markdownTableSplit = (props: SplitProps): SplitResponse => {
let { text = '', chunkSize, maxSize = defaultMaxChunkSize } = props; const { text = '', chunkSize, maxSize = defaultMaxChunkSize } = props;
// split by rows // split by rows
const splitText2Lines = text.split('\n').filter((line) => line.trim()); const splitText2Lines = text.split('\n').filter((line) => line.trim());
...@@ -127,8 +127,8 @@ ${mdSplitString} ...@@ -127,8 +127,8 @@ ${mdSplitString}
5. 标点分割：重叠 5. 标点分割：重叠
*/ */
const commonSplit = (props: SplitProps): SplitResponse => { const commonSplit = (props: SplitProps): SplitResponse => {
let { const {
text = '', text: rawText = '',
chunkSize, chunkSize,
paragraphChunkDeep = 5, paragraphChunkDeep = 5,
paragraphChunkMinSize = 100, paragraphChunkMinSize = 100,
...@@ -136,10 +136,15 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -136,10 +136,15 @@ const commonSplit = (props: SplitProps): SplitResponse => {
overlapRatio = 0.15, overlapRatio = 0.15,
customReg = [] customReg = []
} = props; } = props;
let text = rawText;
const splitMarker = 'SPLIT_HERE_SPLIT_HERE'; const splitMarker = 'SPLIT_HERE_SPLIT_HERE';
const codeBlockMarker = 'CODE_BLOCK_LINE_MARKER'; const codeBlockMarker = 'CODE_BLOCK_LINE_MARKER';
const overlapLen = Math.round(chunkSize * overlapRatio); const overlapLen = Math.round(chunkSize * overlapRatio);
// 代码块需要尽量保留完整性，但不能直接使用模型 maxSize，否则大段正文包在 ```json/markdown``` 中会绕过 chunkSize 形成超大分块。
const maxCodeBlockChunks = 4;
const codeBlockMaxLen = Math.min(maxSize, chunkSize * maxCodeBlockChunks);
const strIsCodeBlock = (str: string) => /^(```[\s\S]*```|~~~[\s\S]*~~~)$/.test(str.trim());
// 特殊模块处理 // 特殊模块处理
// 1. 代码块处理 - 去除空字符 // 1. 代码块处理 - 去除空字符
...@@ -173,14 +178,15 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -173,14 +178,15 @@ const commonSplit = (props: SplitProps): SplitResponse => {
return rules; return rules;
})(paragraphChunkDeep); })(paragraphChunkDeep);
const stepReges: { reg: RegExp | string; maxLen: number }[] = [ const stepReges: { reg: RegExp | string; maxLen: number; splitAround?: boolean }[] = [
...customReg.map((text) => ({ ...customReg.map((text) => ({
reg: text.replace(/\\n/g, '\n'), reg: text.replace(/\\n/g, '\n'),
maxLen: maxSize maxLen: maxSize
})), })),
...markdownHeaderRules, ...markdownHeaderRules,
{ reg: /([\n](```[\s\S]*?```|~~~[\s\S]*?~~~))/g, maxLen: maxSize }, // code block // 代码块需要独立成段，避免吞掉前面大段正文；短代码块仍尽量保持完整。
{ reg: /(^|\n)(```[\s\S]*?```|~~~[\s\S]*?~~~)/g, maxLen: codeBlockMaxLen, splitAround: true },
// HTML Table tag 尽可能保障完整 // HTML Table tag 尽可能保障完整
{ {
reg: /(\n\|(?:[^\n|]*\|)+\n\|(?:[:\-\s]*\|)+\n(?:\|(?:[^\n|]*\|)*\n)*)/g, reg: /(\n\|(?:[^\n|]*\|)+\n\|(?:[:\-\s]*\|)+\n(?:\|(?:[^\n|]*\|)*\n)*)/g,
...@@ -216,7 +222,7 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -216,7 +222,7 @@ const commonSplit = (props: SplitProps): SplitResponse => {
const isCustomStep = checkIsCustomStep(step); const isCustomStep = checkIsCustomStep(step);
const isMarkdownSplit = checkIsMarkdownSplit(step); const isMarkdownSplit = checkIsMarkdownSplit(step);
const { reg, maxLen } = stepReges[step]; const { reg, maxLen, splitAround } = stepReges[step];
const replaceText = (() => { const replaceText = (() => {
if (typeof reg === 'string') { if (typeof reg === 'string') {
...@@ -239,6 +245,7 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -239,6 +245,7 @@ const commonSplit = (props: SplitProps): SplitResponse => {
(() => { (() => {
if (isCustomStep) return splitMarker; if (isCustomStep) return splitMarker;
if (isMarkdownSplit) return `${splitMarker}$1`; if (isMarkdownSplit) return `${splitMarker}$1`;
if (splitAround) return `${splitMarker}$&${splitMarker}`;
return `$1${splitMarker}`; return `$1${splitMarker}`;
})() })()
); );
...@@ -341,6 +348,25 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -341,6 +348,25 @@ const commonSplit = (props: SplitProps): SplitResponse => {
const newText = lastText + currentText; const newText = lastText + currentText;
const newTextLen = getTextValidLength(newText); const newTextLen = getTextValidLength(newText);
// 代码块独立处理，避免“前面正文 + 代码块”被 maxSize 合成超大分块。
if (strIsCodeBlock(currentText)) {
if (lastTextLen > 0) {
chunks.push(lastText);
lastText = '';
}
if (getTextValidLength(currentText) > maxLen) {
const restoredCodeBlock = currentText.replaceAll(codeBlockMarker, '\n');
for (let i = 0; i < restoredCodeBlock.length; i += chunkSize) {
chunks.push(restoredCodeBlock.slice(i, i + chunkSize));
}
} else {
chunks.push(currentText);
}
continue;
}
// split the current table if it will exceed after adding // split the current table if it will exceed after adding
if (strIsMdTable(currentText) && newTextLen > maxLen) { if (strIsMdTable(currentText) && newTextLen > maxLen) {
if (lastTextLen > 0) { if (lastTextLen > 0) {
...@@ -447,7 +473,10 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -447,7 +473,10 @@ const commonSplit = (props: SplitProps): SplitResponse => {
/* If the last chunk is independent, it needs to be push chunks. */ /* If the last chunk is independent, it needs to be push chunks. */
if (lastText && chunks[chunks.length - 1] && !chunks[chunks.length - 1].endsWith(lastText)) { if (lastText && chunks[chunks.length - 1] && !chunks[chunks.length - 1].endsWith(lastText)) {
if (getTextValidLength(lastText) < chunkSize * 0.4) { if (
getTextValidLength(lastText) < chunkSize * 0.4 &&
!strIsCodeBlock(chunks[chunks.length - 1])
) {
chunks[chunks.length - 1] = chunks[chunks.length - 1] + lastText; chunks[chunks.length - 1] = chunks[chunks.length - 1] + lastText;
} else { } else {
chunks.push(lastText); chunks.push(lastText);
...@@ -487,7 +516,7 @@ const commonSplit = (props: SplitProps): SplitResponse => { ...@@ -487,7 +516,7 @@ const commonSplit = (props: SplitProps): SplitResponse => {
* markdown * markdown
*/ */
export const splitText2Chunks = (props: SplitProps): SplitResponse => { export const splitText2Chunks = (props: SplitProps): SplitResponse => {
let { text = '' } = props; const { text = '' } = props;
const splitWithCustomSign = text.split(CUSTOM_SPLIT_SIGN); const splitWithCustomSign = text.split(CUSTOM_SPLIT_SIGN);
const splitResult = splitWithCustomSign.map((item) => { const splitResult = splitWithCustomSign.map((item) => {
......
packages/global/test/common/string/textSplitter.test.ts
...@@ -5,6 +5,7 @@ import fs from 'fs'; ...@@ -5,6 +5,7 @@ import fs from 'fs';
const simpleChunks = (chunks: string[]) => { const simpleChunks = (chunks: string[]) => {
return chunks.map((chunk) => chunk.replace(/\s+/g, '')); return chunks.map((chunk) => chunk.replace(/\s+/g, ''));
}; };
const getValidLength = (text: string) => text.replaceAll(/[\s\n]/g, '').length;
// 简单的嵌套测试 // 简单的嵌套测试
it(`Test splitText2Chunks 1`, () => { it(`Test splitText2Chunks 1`, () => {
...@@ -581,10 +582,36 @@ FastGPT AI 相关参数配置说明 ...@@ -581,10 +582,36 @@ FastGPT AI 相关参数配置说明
maxSize: 100000 maxSize: 100000
}); });
const normalizedChunks = simpleChunks(chunks); expect(chunks[0]).toContain('这是一个测试的内容，包含代码块');
const normalizedExpected = simpleChunks(mock.result); expect(chunks[0]).not.toContain('~~~js');
expect(chunks.some((chunk) => chunk.startsWith('~~~js'))).toBe(true);
expect(chunks.some((chunk) => chunk.endsWith('~~~'))).toBe(true);
expect(chunks.join('\n')).not.toContain('CODE_BLOCK_LINE_MARKER');
expect(chunks[chunks.length - 1]).toContain('最大上下文');
expect(Math.max(...chunks.map(getValidLength))).toBeLessThanOrEqual(500 * 1.2);
});
expect(normalizedChunks).toEqual(normalizedExpected); it(`Test splitText2Chunks 8.1 - code block should not swallow long previous text`, () => {
const longText = Array.from(
{ length: 120 },
(_, index) =>
`第${index}段内容。FastGPT 知识库分块需要按照用户配置的 chunkSize 稳定切分，不能因为后面出现代码块就把前文全部合并。`
).join('\n\n');
const imageCodeBlock = '```markdown![](dataset/xxx.png)```';
const { chunks } = splitText2Chunks({
text: `${longText}\n\n${imageCodeBlock}\n\n后续正文。`,
chunkSize: 1000,
maxSize: 128000,
overlapRatio: 0
});
const imageChunkIndex = chunks.findIndex((chunk) => chunk === imageCodeBlock);
expect(imageChunkIndex).toBeGreaterThan(0);
expect(chunks[imageChunkIndex - 1]).not.toContain(imageCodeBlock);
expect(chunks[imageChunkIndex]).not.toContain('第0段内容');
expect(Math.max(...chunks.map(getValidLength))).toBeLessThanOrEqual(1200);
}); });
// 表格分割测试 - 不超出maxSize // 表格分割测试 - 不超出maxSize
......
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