update textsplitter

textsplitter/__init__.py

@@ -1 +1,2 @@
 from .chinese_text_splitter import ChineseTextSplitter
+from .ali_text_splitter import AliTextSplitter

textsplitter/ali_text_splitter.py

@@ -0,0 +1,26 @@
+from langchain.text_splitter import CharacterTextSplitter
+import re
+from typing import List
+from modelscope.pipelines import pipeline
+
+p = pipeline(
+    task="document-segmentation",
+    model='damo/nlp_bert_document-segmentation_chinese-base',
+    device="cpu")
+
+class AliTextSplitter(CharacterTextSplitter):
+    def __init__(self, pdf: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        self.pdf = pdf
+
+    def split_text(self, text: str) -> List[str]:
+        # use_document_segmentation参数指定是否用语义切分文档，此处采取的文档语义分割模型为达摩院开源的nlp_bert_document-segmentation_chinese-base，论文见https://arxiv.org/abs/2107.09278
+        # 如果使用模型进行文档语义切分，那么需要安装modelscope[nlp]：pip install "modelscope[nlp]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
+        # 考虑到使用了三个模型，可能对于低配置gpu不太友好，因此这里将模型load进cpu计算，有需要的话可以替换device为自己的显卡id
+        if self.pdf:
+            text = re.sub(r"\n{3,}", r"\n", text)
+            text = re.sub('\s', " ", text)
+            text = re.sub("\n\n", "", text)
+        result = p(documents=text)
+        sent_list = [i for i in result["text"].split("\n\t") if i]
+        return sent_list

textsplitter/chinese_text_splitter.py

@@ -9,24 +9,11 @@ class ChineseTextSplitter(CharacterTextSplitter):
         super().__init__(**kwargs)
         self.pdf = pdf
 
-    def split_text1(self, text: str, use_document_segmentation: bool = False) -> List[str]:
-        # use_document_segmentation参数指定是否用语义切分文档，此处采取的文档语义分割模型为达摩院开源的nlp_bert_document-segmentation_chinese-base，论文见https://arxiv.org/abs/2107.09278
-        # 如果使用模型进行文档语义切分，那么需要安装modelscope[nlp]：pip install "modelscope[nlp]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
-        # 考虑到使用了三个模型，可能对于低配置gpu不太友好，因此这里将模型load进cpu计算，有需要的话可以替换device为自己的显卡id
+    def split_text1(self, text: str) -> List[str]:
         if self.pdf:
             text = re.sub(r"\n{3,}", "\n", text)
             text = re.sub('\s', ' ', text)
             text = text.replace("\n\n", "")
-        if use_document_segmentation:
-            from modelscope.pipelines import pipeline
-            p = pipeline(
-                task="document-segmentation",
-                model='damo/nlp_bert_document-segmentation_chinese-base',
-                device="cpu")
-            result = p(documents=text)
-            sent_list = [i for i in result["text"].split("\n\t") if i]
-            return sent_list
-        else:
         sent_sep_pattern = re.compile('([﹒﹔﹖﹗．。！？]["’”」』]{0,2}|(?=["‘“「『]{1,2}|$))')  # del ：；
         sent_list = []
         for ele in sent_sep_pattern.split(text):
@@ -36,21 +23,12 @@ class ChineseTextSplitter(CharacterTextSplitter):
                 sent_list.append(ele)
         return sent_list
 
-    def split_text(self, text: str, use_document_segmentation: bool = False) -> List[str]:
+    def split_text(self, text: str) -> List[str]:
         if self.pdf:
             text = re.sub(r"\n{3,}", r"\n", text)
             text = re.sub('\s', " ", text)
             text = re.sub("\n\n", "", text)
-        if use_document_segmentation:
-            from modelscope.pipelines import pipeline
-            p = pipeline(
-                task="document-segmentation",
-                model='damo/nlp_bert_document-segmentation_chinese-base',
-                device="cpu")
-            result = p(documents=text)
-            sent_list = [i for i in result["text"].split("\n\t") if i]
-            return sent_list
-        else:
+
         text = re.sub(r'([;；.!?。！？\?])([^”’])', r"\1\n\2", text)  # 单字符断句符
         text = re.sub(r'(\.{6})([^"’”」』])', r"\1\n\2", text)  # 英文省略号
         text = re.sub(r'(\…{2})([^"’”」』])', r"\1\n\2", text)  # 中文省略号