package cn.iocoder.yudao.module.ai.service.knowledge.splitter;
|
|
import cn.hutool.core.collection.CollUtil;
|
import cn.hutool.core.util.StrUtil;
|
import lombok.AllArgsConstructor;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.transformer.splitter.TextSplitter;
|
import java.util.ArrayList;
|
import java.util.Collections;
|
import java.util.List;
|
import java.util.regex.Matcher;
|
import java.util.regex.Pattern;
|
|
/**
|
* Markdown QA 格式专用切片器
|
* 识别二级标题（## ）作为问题标记，保持问答对完整性
|
*/
|
@Slf4j
|
public class MarkdownQaSplitter extends TextSplitter {
|
|
private static final Pattern H2_PATTERN = Pattern.compile("^##\\s+(.+)$", Pattern.MULTILINE);
|
private static final String PARAGRAPH_SEPARATOR = "\n\n";
|
private static final Pattern SENTENCE_PATTERN = Pattern.compile("[。！？.!?]\\s*");
|
|
private final int chunkSize;
|
private final TokenEstimator tokenEstimator;
|
|
public MarkdownQaSplitter(int chunkSize) {
|
this.chunkSize = chunkSize;
|
this.tokenEstimator = new SimpleTokenEstimator();
|
}
|
|
@Override
|
protected List<String> splitText(String text) {
|
if (StrUtil.isEmpty(text)) return Collections.emptyList();
|
List<QaPair> qaPairs = parseQaPairs(text);
|
if (CollUtil.isEmpty(qaPairs)) return fallbackSplit(text);
|
List<String> result = new ArrayList<>();
|
for (QaPair qaPair : qaPairs) result.addAll(splitQaPair(qaPair));
|
return result;
|
}
|
|
private List<QaPair> parseQaPairs(String content) {
|
List<QaPair> qaPairs = new ArrayList<>();
|
List<Integer> headingPositions = new ArrayList<>();
|
List<String> questions = new ArrayList<>();
|
Matcher matcher = H2_PATTERN.matcher(content);
|
while (matcher.find()) { headingPositions.add(matcher.start()); questions.add(matcher.group(1).trim()); }
|
if (CollUtil.isEmpty(headingPositions)) return qaPairs;
|
for (int i = 0; i < headingPositions.size(); i++) {
|
int start = headingPositions.get(i);
|
int end = (i + 1 < headingPositions.size()) ? headingPositions.get(i + 1) : content.length();
|
String fullText = content.substring(start, end).trim();
|
String question = questions.get(i);
|
String answer = fullText.substring(fullText.indexOf('\n') + 1).trim();
|
qaPairs.add(new QaPair(question, answer, fullText));
|
}
|
return qaPairs;
|
}
|
|
private List<String> splitQaPair(QaPair qaPair) {
|
int qaTokens = tokenEstimator.estimate(qaPair.fullText);
|
if (qaTokens <= chunkSize) return Collections.singletonList(qaPair.fullText);
|
List<String> chunks = new ArrayList<>();
|
List<String> answerChunks = splitLongAnswer(qaPair.answer, qaPair.question);
|
for (String answerChunk : answerChunks) {
|
chunks.add("## " + qaPair.question + "\n" + answerChunk);
|
}
|
return chunks;
|
}
|
|
private List<String> splitLongAnswer(String answer, String question) {
|
List<String> chunks = new ArrayList<>();
|
String questionHeader = "## " + question + "\n";
|
int questionTokens = tokenEstimator.estimate(questionHeader);
|
int availableTokens = chunkSize - questionTokens - 10;
|
String[] paragraphs = answer.split(PARAGRAPH_SEPARATOR);
|
StringBuilder currentChunk = new StringBuilder();
|
int currentTokens = 0;
|
for (String paragraph : paragraphs) {
|
if (StrUtil.isEmpty(paragraph)) continue;
|
int paragraphTokens = tokenEstimator.estimate(paragraph);
|
if (paragraphTokens > availableTokens) {
|
if (currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
|
chunks.addAll(splitLongParagraph(paragraph, availableTokens));
|
continue;
|
}
|
if (currentTokens + paragraphTokens > availableTokens && currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
|
if (currentChunk.length() > 0) currentChunk.append("\n\n");
|
currentChunk.append(paragraph);
|
currentTokens += paragraphTokens;
|
}
|
if (currentChunk.length() > 0) chunks.add(currentChunk.toString().trim());
|
return CollUtil.isEmpty(chunks) ? Collections.singletonList(answer) : chunks;
|
}
|
|
private List<String> splitLongParagraph(String paragraph, int availableTokens) {
|
List<String> chunks = new ArrayList<>();
|
String[] sentences = SENTENCE_PATTERN.split(paragraph);
|
StringBuilder currentChunk = new StringBuilder();
|
int currentTokens = 0;
|
for (String sentence : sentences) {
|
if (StrUtil.isEmpty(sentence)) continue;
|
int sentenceTokens = tokenEstimator.estimate(sentence);
|
if (sentenceTokens > availableTokens) { if (currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; } chunks.add(sentence.trim()); continue; }
|
if (currentTokens + sentenceTokens > availableTokens && currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
|
currentChunk.append(sentence);
|
currentTokens += sentenceTokens;
|
}
|
if (currentChunk.length() > 0) chunks.add(currentChunk.toString().trim());
|
return chunks.isEmpty() ? Collections.singletonList(paragraph) : chunks;
|
}
|
|
private List<String> fallbackSplit(String content) {
|
List<String> chunks = new ArrayList<>();
|
String[] paragraphs = content.split(PARAGRAPH_SEPARATOR);
|
StringBuilder currentChunk = new StringBuilder();
|
int currentTokens = 0;
|
for (String paragraph : paragraphs) {
|
if (StrUtil.isEmpty(paragraph)) continue;
|
int paragraphTokens = tokenEstimator.estimate(paragraph);
|
if (currentTokens + paragraphTokens > chunkSize && currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
|
if (currentChunk.length() > 0) currentChunk.append("\n\n");
|
currentChunk.append(paragraph);
|
currentTokens += paragraphTokens;
|
}
|
if (currentChunk.length() > 0) chunks.add(currentChunk.toString().trim());
|
return chunks.isEmpty() ? Collections.singletonList(content) : chunks;
|
}
|
|
@AllArgsConstructor
|
private static class QaPair { String question; String answer; String fullText; }
|
|
public interface TokenEstimator { int estimate(String text); }
|
|
private static class SimpleTokenEstimator implements TokenEstimator {
|
@Override
|
public int estimate(String text) {
|
if (StrUtil.isEmpty(text)) return 0;
|
int chineseChars = 0, englishWords = 0;
|
for (char c : text.toCharArray()) { if (c >= 0x4E00 && c <= 0x9FA5) chineseChars++; }
|
for (String word : text.split("\\s+")) { if (word.matches(".*[a-zA-Z].*")) englishWords++; }
|
return chineseChars + (int)(englishWords * 1.3);
|
}
|
}
|
}
|
