package cn.iocoder.yudao.module.ai.service.knowledge.splitter;
|
|
import cn.hutool.core.util.StrUtil;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.transformer.splitter.TextSplitter;
|
import java.util.ArrayList;
|
import java.util.Arrays;
|
import java.util.Collections;
|
import java.util.List;
|
import java.util.regex.Matcher;
|
import java.util.regex.Pattern;
|
|
/**
|
* 语义化文本切片器
|
* 优先在段落边界处切分，其次在句子边界处切分，保持语义完整性
|
*/
|
@Slf4j
|
public class SemanticTextSplitter extends TextSplitter {
|
|
private final int chunkSize;
|
private final int chunkOverlap;
|
private static final List<String> PARAGRAPH_SEPARATORS = Arrays.asList("\n\n\n", "\n\n", "\n");
|
private static final Pattern SENTENCE_END_PATTERN = Pattern.compile("[。！？.!?]+[\\s\"'）)】\\]]*");
|
private final MarkdownQaSplitter.TokenEstimator tokenEstimator;
|
|
public SemanticTextSplitter(int chunkSize, int chunkOverlap) {
|
this.chunkSize = chunkSize;
|
this.chunkOverlap = Math.min(chunkOverlap, chunkSize / 2);
|
this.tokenEstimator = new SimpleTokenEstimator();
|
}
|
|
public SemanticTextSplitter(int chunkSize) { this(chunkSize, 50); }
|
|
@Override
|
protected List<String> splitText(String text) {
|
if (StrUtil.isEmpty(text)) return Collections.emptyList();
|
return splitTextRecursive(text);
|
}
|
|
private List<String> splitTextRecursive(String text) {
|
int textTokens = tokenEstimator.estimate(text);
|
if (textTokens <= chunkSize) return Collections.singletonList(text.trim());
|
|
List<String> splits = null;
|
String usedSeparator = null;
|
for (String separator : PARAGRAPH_SEPARATORS) {
|
if (text.contains(separator)) { splits = Arrays.asList(text.split(Pattern.quote(separator))); usedSeparator = separator; break; }
|
}
|
if (splits == null || splits.size() == 1) { splits = splitBySentences(text); usedSeparator = ""; }
|
return mergeSplits(splits, usedSeparator);
|
}
|
|
private List<String> splitBySentences(String text) {
|
List<String> sentences = new ArrayList<>();
|
int lastEnd = 0;
|
Matcher matcher = SENTENCE_END_PATTERN.matcher(text);
|
while (matcher.find()) {
|
String sentence = text.substring(lastEnd, matcher.end()).trim();
|
if (StrUtil.isNotEmpty(sentence)) sentences.add(sentence);
|
lastEnd = matcher.end();
|
}
|
if (lastEnd < text.length()) { String remaining = text.substring(lastEnd).trim(); if (StrUtil.isNotEmpty(remaining)) sentences.add(remaining); }
|
return sentences.isEmpty() ? Collections.singletonList(text) : sentences;
|
}
|
|
private List<String> mergeSplits(List<String> splits, String separator) {
|
List<String> chunks = new ArrayList<>();
|
List<String> currentChunks = new ArrayList<>();
|
int currentLength = 0;
|
for (String split : splits) {
|
if (StrUtil.isEmpty(split)) continue;
|
int splitTokens = tokenEstimator.estimate(split);
|
if (splitTokens > chunkSize) {
|
if (!currentChunks.isEmpty()) { chunks.add(String.join(separator, currentChunks).trim()); currentChunks.clear(); currentLength = 0; }
|
if (!separator.isEmpty()) { chunks.addAll(splitTextRecursive(split)); } else { chunks.addAll(forceSplitLongText(split)); }
|
continue;
|
}
|
int separatorTokens = StrUtil.isEmpty(separator) ? 0 : tokenEstimator.estimate(separator);
|
if (!currentChunks.isEmpty() && currentLength + splitTokens + separatorTokens > chunkSize) {
|
chunks.add(String.join(separator, currentChunks).trim());
|
currentChunks = getOverlappingChunks(currentChunks, separator);
|
currentLength = estimateTokens(currentChunks, separator);
|
}
|
currentChunks.add(split);
|
currentLength += splitTokens + separatorTokens;
|
}
|
if (!currentChunks.isEmpty()) chunks.add(String.join(separator, currentChunks).trim());
|
return chunks;
|
}
|
|
private List<String> getOverlappingChunks(List<String> chunks, String separator) {
|
if (chunkOverlap == 0 || chunks.isEmpty()) return new ArrayList<>();
|
List<String> overlapping = new ArrayList<>();
|
int tokens = 0;
|
for (int i = chunks.size() - 1; i >= 0; i--) {
|
String chunk = chunks.get(i);
|
int chunkTokens = tokenEstimator.estimate(chunk);
|
if (tokens + chunkTokens > chunkOverlap) break;
|
overlapping.add(0, chunk);
|
tokens += chunkTokens + (StrUtil.isEmpty(separator) ? 0 : tokenEstimator.estimate(separator));
|
}
|
return overlapping;
|
}
|
|
private int estimateTokens(List<String> chunks, String separator) {
|
int total = 0;
|
for (int i = 0; i < chunks.size(); i++) { total += tokenEstimator.estimate(chunks.get(i)); if (i < chunks.size() - 1 && StrUtil.isNotEmpty(separator)) total += tokenEstimator.estimate(separator); }
|
return total;
|
}
|
|
private List<String> forceSplitLongText(String text) {
|
List<String> chunks = new ArrayList<>();
|
int charsPerChunk = (int)(chunkSize * 0.8);
|
for (int i = 0; i < text.length(); i += charsPerChunk) { int end = Math.min(i + charsPerChunk, text.length()); chunks.add(text.substring(i, end).trim()); }
|
return chunks;
|
}
|
|
private static class SimpleTokenEstimator implements MarkdownQaSplitter.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);
|
}
|
}
|
}
|
