80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
import re
|
|
|
|
|
|
# rule based segmentation based on https://stackoverflow.com/a/31505798, works surprisingly well
|
|
def split_sentences(
|
|
text: str, min_sentence_len: int = 20, retain_format: bool = False
|
|
) -> list[tuple[str, int, int]]:
|
|
"""
|
|
the text may not contain substrings "<prd>" or "<stop>"
|
|
"""
|
|
alphabets = r"([A-Za-z])"
|
|
prefixes = r"(Mr|St|Mrs|Ms|Dr)[.]"
|
|
suffixes = r"(Inc|Ltd|Jr|Sr|Co)"
|
|
starters = r"(Mr|Mrs|Ms|Dr|Prof|Capt|Cpt|Lt|He\s|She\s|It\s|They\s|Their\s|Our\s|We\s|But\s|However\s|That\s|This\s|Wherever)" # noqa: E501
|
|
acronyms = r"([A-Z][.][A-Z][.](?:[A-Z][.])?)"
|
|
websites = r"[.](com|net|org|io|gov|edu|me)"
|
|
digits = r"([0-9])"
|
|
multiple_dots = r"\.{2,}"
|
|
|
|
# fmt: off
|
|
if retain_format:
|
|
text = text.replace("\n","<nel><stop>")
|
|
else:
|
|
text = text.replace("\n"," ")
|
|
|
|
text = re.sub(prefixes,"\\1<prd>", text)
|
|
text = re.sub(websites,"<prd>\\1", text)
|
|
text = re.sub(digits + "[.]" + digits,"\\1<prd>\\2",text)
|
|
# text = re.sub(multiple_dots, lambda match: "<prd>" * len(match.group(0)) + "<stop>", text)
|
|
# TODO(theomonnom): need improvement for ""..." dots", check capital + next sentence should not be # noqa: E501
|
|
# small
|
|
text = re.sub(multiple_dots, lambda match: "<prd>" * len(match.group(0)), text)
|
|
if "Ph.D" in text:
|
|
text = text.replace("Ph.D.","Ph<prd>D<prd>")
|
|
text = re.sub(r"\s" + alphabets + "[.] "," \\1<prd> ",text)
|
|
text = re.sub(acronyms+" "+starters,"\\1<stop> \\2",text)
|
|
text = re.sub(alphabets + "[.]" + alphabets + "[.]" + alphabets + "[.]","\\1<prd>\\2<prd>\\3<prd>",text) # noqa: E501
|
|
text = re.sub(alphabets + "[.]" + alphabets + "[.]","\\1<prd>\\2<prd>",text)
|
|
text = re.sub(r" "+suffixes+"[.] "+starters," \\1<stop> \\2",text)
|
|
text = re.sub(r" "+suffixes+"[.]"," \\1<prd>",text)
|
|
text = re.sub(r" " + alphabets + "[.]"," \\1<prd>",text)
|
|
|
|
# mark end of sentence punctuations with <stop>
|
|
text = re.sub(r"([.!?。!?])([\"”])", "\\1\\2<stop>", text)
|
|
text = re.sub(r"([.!?。!?])(?![\"”])", "\\1<stop>", text)
|
|
|
|
text = text.replace("<prd>",".")
|
|
# fmt: on
|
|
|
|
if retain_format:
|
|
text = text.replace("<nel>", "\n")
|
|
splitted_sentences = text.split("<stop>")
|
|
text = text.replace("<stop>", "")
|
|
|
|
sentences: list[tuple[str, int, int]] = []
|
|
|
|
buff = ""
|
|
start_pos = 0
|
|
end_pos = 0
|
|
pre_pad = "" if retain_format else " "
|
|
for match in splitted_sentences:
|
|
if retain_format:
|
|
sentence = match
|
|
else:
|
|
sentence = match.strip()
|
|
if not sentence:
|
|
continue
|
|
|
|
buff += pre_pad + sentence
|
|
end_pos += len(match)
|
|
if len(buff) > min_sentence_len:
|
|
sentences.append((buff[len(pre_pad) :], start_pos, end_pos))
|
|
start_pos = end_pos
|
|
buff = ""
|
|
|
|
if buff:
|
|
sentences.append((buff[len(pre_pad) :], start_pos, len(text) - 1))
|
|
|
|
return sentences
|