100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
from langchain.docstore.document import Document
|
||
import re
|
||
|
||
|
||
def under_non_alpha_ratio(text: str, threshold: float = 0.5):
|
||
"""Checks if the proportion of non-alpha characters in the text snippet exceeds a given
|
||
threshold. This helps prevent text like "-----------BREAK---------" from being tagged
|
||
as a title or narrative text. The ratio does not count spaces.
|
||
|
||
Parameters
|
||
----------
|
||
text
|
||
The input string to test
|
||
threshold
|
||
If the proportion of non-alpha characters exceeds this threshold, the function
|
||
returns False
|
||
"""
|
||
if len(text) == 0:
|
||
return False
|
||
|
||
alpha_count = len([char for char in text if char.strip() and char.isalpha()])
|
||
total_count = len([char for char in text if char.strip()])
|
||
try:
|
||
ratio = alpha_count / total_count
|
||
return ratio < threshold
|
||
except:
|
||
return False
|
||
|
||
|
||
def is_possible_title(
|
||
text: str,
|
||
title_max_word_length: int = 20,
|
||
non_alpha_threshold: float = 0.5,
|
||
) -> bool:
|
||
"""Checks to see if the text passes all of the checks for a valid title.
|
||
|
||
Parameters
|
||
----------
|
||
text
|
||
The input text to check
|
||
title_max_word_length
|
||
The maximum number of words a title can contain
|
||
non_alpha_threshold
|
||
The minimum number of alpha characters the text needs to be considered a title
|
||
"""
|
||
|
||
# 文本长度为0的话，肯定不是title
|
||
if len(text) == 0:
|
||
print("Not a title. Text is empty.")
|
||
return False
|
||
|
||
# 文本中有标点符号，就不是title
|
||
ENDS_IN_PUNCT_PATTERN = r"[^\w\s]\Z"
|
||
ENDS_IN_PUNCT_RE = re.compile(ENDS_IN_PUNCT_PATTERN)
|
||
if ENDS_IN_PUNCT_RE.search(text) is not None:
|
||
return False
|
||
|
||
# 文本长度不能超过设定值，默认20
|
||
# NOTE(robinson) - splitting on spaces here instead of word tokenizing because it
|
||
# is less expensive and actual tokenization doesn't add much value for the length check
|
||
if len(text) > title_max_word_length:
|
||
return False
|
||
|
||
# 文本中数字的占比不能太高，否则不是title
|
||
if under_non_alpha_ratio(text, threshold=non_alpha_threshold):
|
||
return False
|
||
|
||
# NOTE(robinson) - Prevent flagging salutations like "To My Dearest Friends," as titles
|
||
if text.endswith((",", ".", "，", "。")):
|
||
return False
|
||
|
||
if text.isnumeric():
|
||
print(f"Not a title. Text is all numeric:\n\n{text}") # type: ignore
|
||
return False
|
||
|
||
# 开头的字符内应该有数字，默认5个字符内
|
||
if len(text) < 5:
|
||
text_5 = text
|
||
else:
|
||
text_5 = text[:5]
|
||
alpha_in_text_5 = sum(list(map(lambda x: x.isnumeric(), list(text_5))))
|
||
if not alpha_in_text_5:
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
def zh_title_enhance(docs: Document) -> Document:
|
||
title = None
|
||
if len(docs) > 0:
|
||
for doc in docs:
|
||
if is_possible_title(doc.page_content):
|
||
doc.metadata['category'] = 'cn_Title'
|
||
title = doc.page_content
|
||
elif title:
|
||
doc.page_content = f"下文与({title})有关。{doc.page_content}"
|
||
return docs
|
||
else:
|
||
print("文件不存在")
|