chinese_text_splitter.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. from langchain.text_splitter import CharacterTextSplitter
  2. import re
  3. from typing import List
  4. from configs.model_config import SENTENCE_SIZE
  5. class ChineseTextSplitter(CharacterTextSplitter):
  6. def __init__(self, pdf: bool = False, **kwargs):
  7. super().__init__(**kwargs)
  8. self.pdf = pdf
  9. def split_text1(self, text: str, use_document_segmentation: bool = False) -> List[str]:
  10. # use_document_segmentation参数指定是否用语义切分文档,此处采取的文档语义分割模型为达摩院开源的nlp_bert_document-segmentation_chinese-base,论文见https://arxiv.org/abs/2107.09278
  11. # 如果使用模型进行文档语义切分,那么需要安装modelscope[nlp]:pip install "modelscope[nlp]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
  12. # 考虑到使用了三个模型,可能对于低配置gpu不太友好,因此这里将模型load进cpu计算,有需要的话可以替换device为自己的显卡id
  13. if self.pdf:
  14. text = re.sub(r"\n{3,}", "\n", text)
  15. text = re.sub('\s', ' ', text)
  16. text = text.replace("\n\n", "")
  17. if use_document_segmentation:
  18. result = p(documents=text)
  19. sent_list = [i for i in result["text"].split("\n\t") if i]
  20. else:
  21. sent_sep_pattern = re.compile('([﹒﹔﹖﹗.。!?]["’”」』]{0,2}|(?=["‘“「『]{1,2}|$))') # del :;
  22. sent_list = []
  23. for ele in sent_sep_pattern.split(text):
  24. if sent_sep_pattern.match(ele) and sent_list:
  25. sent_list[-1] += ele
  26. elif ele:
  27. sent_list.append(ele)
  28. return sent_list
  29. def split_text(self, text: str, use_document_segmentation: bool = False) -> List[str]:
  30. if self.pdf:
  31. text = re.sub(r"\n{3,}", r"\n", text)
  32. text = re.sub('\s', " ", text)
  33. text = re.sub("\n\n", "", text)
  34. if use_document_segmentation:
  35. from modelscope.pipelines import pipeline
  36. p = pipeline(
  37. task="document-segmentation",
  38. model='damo/nlp_bert_document-segmentation_chinese-base',
  39. device="cpu")
  40. result = p(documents=text)
  41. sent_list = [i for i in result["text"].split("\n\t") if i]
  42. return sent_list
  43. else:
  44. text = re.sub(r'([;;.!?。!?\?])([^”’])', r"\1\n\2", text) # 单字符断句符
  45. text = re.sub(r'(\.{6})([^"’”」』])', r"\1\n\2", text) # 英文省略号
  46. text = re.sub(r'(\…{2})([^"’”」』])', r"\1\n\2", text) # 中文省略号
  47. text = re.sub(r'([;;!?。!?\?]["’”」』]{0,2})([^;;!?,。!?\?])', r'\1\n\2', text)
  48. # 如果双引号前有终止符,那么双引号才是句子的终点,把分句符\n放到双引号后,注意前面的几句都小心保留了双引号
  49. text = text.rstrip() # 段尾如果有多余的\n就去掉它
  50. # 很多规则中会考虑分号;,但是这里我把它忽略不计,破折号、英文双引号等同样忽略,需要的再做些简单调整即可。
  51. ls = [i for i in text.split("\n") if i]
  52. for ele in ls:
  53. if len(ele) > SENTENCE_SIZE:
  54. ele1 = re.sub(r'([,,.]["’”」』]{0,2})([^,,.])', r'\1\n\2', ele)
  55. ele1_ls = ele1.split("\n")
  56. for ele_ele1 in ele1_ls:
  57. if len(ele_ele1) > SENTENCE_SIZE:
  58. ele_ele2 = re.sub(r'([\n]{1,}| {2,}["’”」』]{0,2})([^\s])', r'\1\n\2', ele_ele1)
  59. ele2_ls = ele_ele2.split("\n")
  60. for ele_ele2 in ele2_ls:
  61. if len(ele_ele2) > SENTENCE_SIZE:
  62. ele_ele3 = re.sub('( ["’”」』]{0,2})([^ ])', r'\1\n\2', ele_ele2)
  63. ele2_id = ele2_ls.index(ele_ele2)
  64. ele2_ls = ele2_ls[:ele2_id] + [i for i in ele_ele3.split("\n") if i] + ele2_ls[
  65. ele2_id + 1:]
  66. ele_id = ele1_ls.index(ele_ele1)
  67. ele1_ls = ele1_ls[:ele_id] + [i for i in ele2_ls if i] + ele1_ls[ele_id + 1:]
  68. id = ls.index(ele)
  69. ls = ls[:id] + [i for i in ele1_ls if i] + ls[id + 1:]
  70. return ls