Langchain-Chatchat
34 строки · 1.6 Кб
1from langchain.text_splitter import CharacterTextSplitter
2import re
3from typing import List
4
5
6class AliTextSplitter(CharacterTextSplitter):
7def __init__(self, pdf: bool = False, **kwargs):
8super().__init__(**kwargs)
9self.pdf = pdf
10
11def split_text(self, text: str) -> List[str]:
12# use_document_segmentation参数指定是否用语义切分文档,此处采取的文档语义分割模型为达摩院开源的nlp_bert_document-segmentation_chinese-base,论文见https://arxiv.org/abs/2107.09278
13# 如果使用模型进行文档语义切分,那么需要安装modelscope[nlp]:pip install "modelscope[nlp]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
14# 考虑到使用了三个模型,可能对于低配置gpu不太友好,因此这里将模型load进cpu计算,有需要的话可以替换device为自己的显卡id
15if self.pdf:
16text = re.sub(r"\n{3,}", r"\n", text)
17text = re.sub('\s', " ", text)
18text = re.sub("\n\n", "", text)
19try:
20from modelscope.pipelines import pipeline
21except ImportError:
22raise ImportError(
23"Could not import modelscope python package. "
24"Please install modelscope with `pip install modelscope`. "
25)
26
27
28p = pipeline(
29task="document-segmentation",
30model='damo/nlp_bert_document-segmentation_chinese-base',
31device="cpu")
32result = p(documents=text)
33sent_list = [i for i in result["text"].split("\n\t") if i]
34return sent_list
35