396 lines
13 KiB
Python
396 lines
13 KiB
Python
|
"""
|
|||
|
|
Semantic Text Splitter using Online Embedding APIs
|
|||
|
|
基于在线 Embedding API 的语义分割器
|
|||
|
|
"""
|
|||
|
|
import re
|
|||
|
|
import asyncio
|
|||
|
|
import httpx
|
|||
|
|
import numpy as np
|
|||
|
|
from typing import List, Dict, Optional
|
|||
|
|
from abc import ABC, abstractmethod
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EmbeddingProvider(ABC):
|
|||
|
|
"""Embedding API 提供商基类"""
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
async def get_embeddings(self, texts: List[str]) -> List[List[float]]:
|
|||
|
|
"""获取文本的嵌入向量"""
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OpenAIEmbedding(EmbeddingProvider):
|
|||
|
|
"""OpenAI 兼容的 Embedding API"""
|
|||
|
|
|
|||
|
|
def __init__(self, api_key: str, base_url: str, model: str = "text-embedding-3-small"):
|
|||
|
|
self.api_key = api_key
|
|||
|
|
self.base_url = base_url.rstrip('/')
|
|||
|
|
self.model = model
|
|||
|
|
|
|||
|
|
async def get_embeddings(self, texts: List[str]) -> List[List[float]]:
|
|||
|
|
"""调用 OpenAI 兼容的 Embedding API"""
|
|||
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|||
|
|
headers = {
|
|||
|
|
"Authorization": f"Bearer {self.api_key}",
|
|||
|
|
"Content-Type": "application/json"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# OpenAI 格式
|
|||
|
|
payload = {
|
|||
|
|
"input": texts,
|
|||
|
|
"model": self.model
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
response = await client.post(
|
|||
|
|
f"{self.base_url}/embeddings",
|
|||
|
|
headers=headers,
|
|||
|
|
json=payload
|
|||
|
|
)
|
|||
|
|
response.raise_for_status()
|
|||
|
|
data = response.json()
|
|||
|
|
|
|||
|
|
# 提取 embeddings
|
|||
|
|
return [item["embedding"] for item in data["data"]]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MiniMaxEmbedding(EmbeddingProvider):
|
|||
|
|
"""MiniMax Embedding API"""
|
|||
|
|
|
|||
|
|
def __init__(self, api_key: str, base_url: str = "https://api.minimax.chat/v1"):
|
|||
|
|
self.api_key = api_key
|
|||
|
|
self.base_url = base_url.rstrip('/')
|
|||
|
|
|
|||
|
|
async def get_embeddings(self, texts: List[str]) -> List[List[float]]:
|
|||
|
|
"""调用 MiniMax Embedding API"""
|
|||
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|||
|
|
headers = {
|
|||
|
|
"Authorization": f"Bearer {self.api_key}",
|
|||
|
|
"Content-Type": "application/json"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# MiniMax 格式
|
|||
|
|
payload = {
|
|||
|
|
"texts": texts,
|
|||
|
|
"model": "embo-01"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
response = await client.post(
|
|||
|
|
f"{self.base_url}/text_embeddings",
|
|||
|
|
headers=headers,
|
|||
|
|
json=payload
|
|||
|
|
)
|
|||
|
|
response.raise_for_status()
|
|||
|
|
data = response.json()
|
|||
|
|
|
|||
|
|
# MiniMax 返回格式可能不同，需要适配
|
|||
|
|
if "data" in data:
|
|||
|
|
return [item["embedding"] for item in data["data"]]
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EmbeddingSplitter:
|
|||
|
|
"""基于 Embedding 的语义分割器基类"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
chunk_size: int = 500,
|
|||
|
|
overlap: int = 50,
|
|||
|
|
embedding_provider: Optional[EmbeddingProvider] = None,
|
|||
|
|
similarity_threshold: float = 0.3,
|
|||
|
|
min_chunk_size: int = 100,
|
|||
|
|
window_size: int = 3
|
|||
|
|
):
|
|||
|
|
self.chunk_size = chunk_size
|
|||
|
|
self.overlap = overlap
|
|||
|
|
self.embedding_provider = embedding_provider
|
|||
|
|
self.similarity_threshold = similarity_threshold
|
|||
|
|
self.min_chunk_size = min_chunk_size
|
|||
|
|
self.window_size = window_size
|
|||
|
|
|
|||
|
|
def _tokenize_sentences(self, text: str) -> List[str]:
|
|||
|
|
"""将文本切分为句子"""
|
|||
|
|
# 中英文句末符号
|
|||
|
|
# 先按换行分割，保持段落结构
|
|||
|
|
paragraphs = re.split(r'\n+', text)
|
|||
|
|
|
|||
|
|
sentences = []
|
|||
|
|
for para in paragraphs:
|
|||
|
|
if not para.strip():
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 按句子符号分割
|
|||
|
|
# 中文：。！？；
|
|||
|
|
# 英文：. ! ? ;
|
|||
|
|
parts = re.split(r'([。！？；\n]|(?<=[.!?])\s+)', para)
|
|||
|
|
|
|||
|
|
# 重新组合句子
|
|||
|
|
current_sentence = ""
|
|||
|
|
for part in parts:
|
|||
|
|
if part in '。！？；.\n':
|
|||
|
|
if current_sentence.strip():
|
|||
|
|
sentences.append(current_sentence.strip())
|
|||
|
|
current_sentence = ""
|
|||
|
|
elif part and part.strip():
|
|||
|
|
current_sentence += part
|
|||
|
|
# 处理最后一个句子
|
|||
|
|
if current_sentence.strip():
|
|||
|
|
sentences.append(current_sentence.strip())
|
|||
|
|
|
|||
|
|
return sentences
|
|||
|
|
|
|||
|
|
def _compute_similarities(self, embeddings: List[List[float]]) -> List[float]:
|
|||
|
|
"""计算相邻句子的余弦相似度"""
|
|||
|
|
similarities = []
|
|||
|
|
|
|||
|
|
for i in range(len(embeddings) - 1):
|
|||
|
|
# 余弦相似度
|
|||
|
|
vec1 = np.array(embeddings[i])
|
|||
|
|
vec2 = np.array(embeddings[i + 1])
|
|||
|
|
|
|||
|
|
# 归一化
|
|||
|
|
vec1 = vec1 / (np.linalg.norm(vec1) + 1e-8)
|
|||
|
|
vec2 = vec2 / (np.linalg.norm(vec2) + 1e-8)
|
|||
|
|
|
|||
|
|
# 点积 = 余弦相似度（归一化后）
|
|||
|
|
sim = np.dot(vec1, vec2)
|
|||
|
|
similarities.append(float(sim))
|
|||
|
|
|
|||
|
|
return similarities
|
|||
|
|
|
|||
|
|
def _smooth_similarities(self, similarities: List[float]) -> List[float]:
|
|||
|
|
"""滑动窗口平滑相似度"""
|
|||
|
|
if not similarities:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
window = self.window_size
|
|||
|
|
smoothed = []
|
|||
|
|
|
|||
|
|
for i in range(len(similarities)):
|
|||
|
|
start = max(0, i - window + 1)
|
|||
|
|
end = i + 1
|
|||
|
|
window_vals = similarities[start:end]
|
|||
|
|
smoothed.append(sum(window_vals) / len(window_vals))
|
|||
|
|
|
|||
|
|
return smoothed
|
|||
|
|
|
|||
|
|
def _detect_boundaries(self, similarities: List[float]) -> List[int]:
|
|||
|
|
"""检测分割点（相似度显著下降的位置）"""
|
|||
|
|
if not similarities:
|
|||
|
|
return [0]
|
|||
|
|
|
|||
|
|
# 平滑
|
|||
|
|
smoothed = self._smooth_similarities(similarities)
|
|||
|
|
|
|||
|
|
# 计算深度分数（类似 TextTiling）
|
|||
|
|
depth_scores = []
|
|||
|
|
for i in range(1, len(smoothed) - 1):
|
|||
|
|
# 当前位置的深度 = 当前位置的值 - 平均值
|
|||
|
|
# 但更准确的是：左侧平均 - 右侧平均
|
|||
|
|
left_avg = sum(smoothed[max(0, i - self.window_size):i]) / self.window_size
|
|||
|
|
right_avg = sum(smoothed[i:min(len(smoothed), i + self.window_size)]) / self.window_size
|
|||
|
|
depth = left_avg - right_avg
|
|||
|
|
depth_scores.append(depth)
|
|||
|
|
|
|||
|
|
# 如果没有足够的点，直接返回
|
|||
|
|
if not depth_scores:
|
|||
|
|
return [0]
|
|||
|
|
|
|||
|
|
# 阈值判断
|
|||
|
|
mean_depth = np.mean(depth_scores)
|
|||
|
|
std_depth = np.std(depth_scores)
|
|||
|
|
|
|||
|
|
# 找分割点：depth 显著高于均值的位置
|
|||
|
|
threshold = mean_depth + 0.5 * std_depth
|
|||
|
|
|
|||
|
|
boundaries = [0] # 起始点
|
|||
|
|
for i, depth in enumerate(depth_scores):
|
|||
|
|
if depth > threshold and depth > self.similarity_threshold:
|
|||
|
|
boundaries.append(i + 1) # 对应相似度的下一个位置
|
|||
|
|
boundaries.append(len(self._tokenize_sentences.__name__)) # 结束点
|
|||
|
|
|
|||
|
|
return sorted(list(set(boundaries)))
|
|||
|
|
|
|||
|
|
def _assemble_chunks(self, sentences: List[str], boundaries: List[int]) -> List[Dict]:
|
|||
|
|
"""按分割点组装 chunks"""
|
|||
|
|
if not sentences:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 重新计算 boundaries（确保不超过句子数）
|
|||
|
|
if not boundaries or boundaries[0] != 0:
|
|||
|
|
boundaries = [0] + boundaries
|
|||
|
|
if boundaries[-1] != len(sentences):
|
|||
|
|
boundaries.append(len(sentences))
|
|||
|
|
|
|||
|
|
chunks = []
|
|||
|
|
for i in range(len(boundaries) - 1):
|
|||
|
|
start = boundaries[i]
|
|||
|
|
end = boundaries[i + 1]
|
|||
|
|
chunk_text = ' '.join(sentences[start:end])
|
|||
|
|
|
|||
|
|
# 如果 chunk 过大，递归分割
|
|||
|
|
if len(chunk_text) > self.chunk_size * 1.5:
|
|||
|
|
# 使用更小的窗口再次分割
|
|||
|
|
sub_chunks = self._split_large_chunk(sentences[start:end])
|
|||
|
|
for j, sub in enumerate(sub_chunks):
|
|||
|
|
chunks.append({
|
|||
|
|
"index": len(chunks),
|
|||
|
|
"content": sub.strip(),
|
|||
|
|
"word_count": len(sub.split()),
|
|||
|
|
"char_count": len(sub)
|
|||
|
|
})
|
|||
|
|
else:
|
|||
|
|
chunks.append({
|
|||
|
|
"index": len(chunks),
|
|||
|
|
"content": chunk_text.strip(),
|
|||
|
|
"word_count": len(chunk_text.split()),
|
|||
|
|
"char_count": len(chunk_text)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 合并过小的相邻 chunks
|
|||
|
|
chunks = self._merge_small_chunks(chunks)
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def _split_large_chunk(self, sentences: List[str]) -> List[str]:
|
|||
|
|
"""分割过大的 chunk"""
|
|||
|
|
# 使用固定长度分割
|
|||
|
|
result = []
|
|||
|
|
current = ""
|
|||
|
|
|
|||
|
|
for sent in sentences:
|
|||
|
|
if len(current) + len(sent) > self.chunk_size:
|
|||
|
|
if current:
|
|||
|
|
result.append(current)
|
|||
|
|
current = sent
|
|||
|
|
else:
|
|||
|
|
current += " " + sent if current else sent
|
|||
|
|
|
|||
|
|
if current:
|
|||
|
|
result.append(current)
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def _merge_small_chunks(self, chunks: List[Dict]) -> List[Dict]:
|
|||
|
|
"""合并过小的相邻 chunks"""
|
|||
|
|
if len(chunks) <= 1:
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
merged = [chunks[0]]
|
|||
|
|
|
|||
|
|
for chunk in chunks[1:]:
|
|||
|
|
# 如果前一个 chunk 太小，合并
|
|||
|
|
if merged[-1]["char_count"] < self.min_chunk_size:
|
|||
|
|
merged[-1]["content"] += " " + chunk["content"]
|
|||
|
|
merged[-1]["word_count"] += chunk["word_count"]
|
|||
|
|
merged[-1]["char_count"] += chunk["char_count"]
|
|||
|
|
else:
|
|||
|
|
merged.append(chunk)
|
|||
|
|
|
|||
|
|
return merged
|
|||
|
|
|
|||
|
|
async def split_with_embedding(self, text: str) -> List[Dict]:
|
|||
|
|
"""使用 Embedding 进行语义分割"""
|
|||
|
|
# 1. 句子切分
|
|||
|
|
sentences = self._tokenize_sentences(text)
|
|||
|
|
if not sentences:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 过滤过短的句子
|
|||
|
|
sentences = [s for s in sentences if len(s) >= 10]
|
|||
|
|
|
|||
|
|
if not sentences:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 2. 如果只有一个句子，直接返回
|
|||
|
|
if len(sentences) == 1:
|
|||
|
|
return [{
|
|||
|
|
"index": 0,
|
|||
|
|
"content": sentences[0],
|
|||
|
|
"word_count": len(sentences[0].split()),
|
|||
|
|
"char_count": len(sentences[0])
|
|||
|
|
}]
|
|||
|
|
|
|||
|
|
# 3. 调用 Embedding API
|
|||
|
|
try:
|
|||
|
|
embeddings = await self.embedding_provider.get_embeddings(sentences)
|
|||
|
|
except Exception as e:
|
|||
|
|
# 如果 embedding 失败，降级到规则分割
|
|||
|
|
print(f"Embedding failed, falling back to rule-based: {e}")
|
|||
|
|
return self._fallback_split(text)
|
|||
|
|
|
|||
|
|
# 4. 计算相似度
|
|||
|
|
similarities = self._compute_similarities(embeddings)
|
|||
|
|
|
|||
|
|
# 5. 检测分割点
|
|||
|
|
boundaries = self._detect_boundaries(similarities)
|
|||
|
|
|
|||
|
|
# 6. 组装 chunks
|
|||
|
|
chunks = self._assemble_chunks(sentences, boundaries)
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def _fallback_split(self, text: str) -> List[Dict]:
|
|||
|
|
"""降级到规则分割"""
|
|||
|
|
# 使用 langchain 的 RecursiveCharacterTextSplitter
|
|||
|
|
splitter = RecursiveCharacterTextSplitter(
|
|||
|
|
chunk_size=self.chunk_size,
|
|||
|
|
chunk_overlap=self.overlap,
|
|||
|
|
separators=["\n\n", "\n", "。", "！", "？", ". ", "! ", "? "]
|
|||
|
|
)
|
|||
|
|
chunks = splitter.split_text(text)
|
|||
|
|
return [{
|
|||
|
|
"index": i,
|
|||
|
|
"content": c.strip(),
|
|||
|
|
"word_count": len(c.split()),
|
|||
|
|
"char_count": len(c)
|
|||
|
|
} for i, c in enumerate(chunks)]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SemanticEmbeddingSplitter(EmbeddingSplitter):
|
|||
|
|
"""基于在线 Embedding 的语义分割器"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
chunk_size: int = 500,
|
|||
|
|
overlap: int = 50,
|
|||
|
|
embedding_provider: Optional[EmbeddingProvider] = None,
|
|||
|
|
similarity_threshold: float = 0.3,
|
|||
|
|
min_chunk_size: int = 100,
|
|||
|
|
window_size: int = 3
|
|||
|
|
):
|
|||
|
|
super().__init__(
|
|||
|
|
chunk_size=chunk_size,
|
|||
|
|
overlap=overlap,
|
|||
|
|
embedding_provider=embedding_provider,
|
|||
|
|
similarity_threshold=similarity_threshold,
|
|||
|
|
min_chunk_size=min_chunk_size,
|
|||
|
|
window_size=window_size
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def split(self, text: str) -> List[Dict]:
|
|||
|
|
"""同步接口，内部调用异步"""
|
|||
|
|
# 由于 split 是同步方法，需要创建新的事件循环
|
|||
|
|
try:
|
|||
|
|
loop = asyncio.get_event_loop()
|
|||
|
|
if loop.is_running():
|
|||
|
|
# 如果在异步环境中，创建新任务
|
|||
|
|
import concurrent.futures
|
|||
|
|
with concurrent.futures.ThreadPoolExecutor() as pool:
|
|||
|
|
future = pool.submit(asyncio.run, self.split_with_embedding(text))
|
|||
|
|
return future.result()
|
|||
|
|
else:
|
|||
|
|
return loop.run_until_complete(self.split_with_embedding(text))
|
|||
|
|
except RuntimeError:
|
|||
|
|
# 没有事件循环，直接创建
|
|||
|
|
return asyncio.run(self.split_with_embedding(text))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def create_embedding_provider(provider: str, api_key: str, base_url: str, model: str = None) -> EmbeddingProvider:
|
|||
|
|
"""创建 Embedding 提供商"""
|
|||
|
|
if provider in ["openai", "compatible"]:
|
|||
|
|
return OpenAIEmbedding(api_key, base_url, model or "text-embedding-3-small")
|
|||
|
|
elif provider == "minimax":
|
|||
|
|
return MiniMaxEmbedding(api_key, base_url)
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"Unsupported embedding provider: {provider}")
|