feat: improve title extraction

- Add support for Chinese and English punctuation
- Fix title truncation and ellipsis
- Improve content cleaning
- Add comprehensive test cases
This commit is contained in:
zhukang 2025-01-15 22:18:55 +08:00
parent 5330de49fc
commit ae32c9a94a
3 changed files with 165 additions and 24 deletions

View File

@ -34,35 +34,77 @@ class DocumentProcessor:
json.dump(self.categories, f, ensure_ascii=False, indent=2)
def extract_title(self, content):
"""从内容中提取或生成标题"""
# 首先尝试从内容的第一行提取标题
lines = content.strip().split('\n')
first_line = lines[0].strip()
"""从内容中提取标题"""
if not content or not content.strip():
return "未命名笔记"
# 清理内容,移除多余的换行符
content = re.sub(r'\n{3,}', '\n\n', content.strip())
# 如果第一行是markdown标题格式直接使用
if first_line.startswith('#'):
return first_line.lstrip('#').strip()
# 尝试从文件名中提取标题
filename_match = re.search(r'captured_text_\d{8}_\d{6}', content)
if filename_match:
content = content.replace(filename_match.group(0), '').strip()
lines = content.split('\n')
first_line = lines[0].strip() if lines else ''
# 如果第一行包含"文章概要"或类似的标记,使用下一个非空行
if any(marker in first_line for marker in ['文章概要', '概要', '总结']):
for line in lines[1:]:
if line.strip():
first_line = line.strip()
break
# 如果第一行以数字和点开头(如 1. 2. 等),尝试使用下一个有效行
if re.match(r'^\d+\.', first_line):
for line in lines[1:]:
line = line.strip()
if line and not re.match(r'^\d+\.', line):
first_line = line
break
# 如果第一行太长,尝试提取第一句话
if len(first_line) > 50:
# 首先尝试按标点符号分割
# 使用正则表达式匹配中英文标点符号
punctuation_pattern = r'([。!?.!?,])'
sentences = re.split(punctuation_pattern, first_line)
# 如果第一行不太长且不是很短,可能是标题
if 10 <= len(first_line) <= 100 and not first_line.endswith('') and not first_line.endswith(':'):
return first_line
# 重新组合第一个句子(包括标点符号)
if len(sentences) > 1:
first_sentence = ''
for i in range(0, min(2, len(sentences))):
first_sentence += sentences[i]
first_line = first_sentence
# 如果第一句话仍然太长,截断它
if len(first_line) > 50:
# 移除末尾的标点符号
first_line = re.sub(r'[。!?.!?,]$', '', first_line[:47]).strip() + '...'
# 尝试查找文章的主题句
for line in lines[:5]: # 只查看前5行
line = line.strip()
if '主题' in line or '概要' in line or '总结' in line:
# 提取冒号或者是后面的内容
if ':' in line or '' in line:
return line.split(':', 1)[1].strip() if ':' in line else line.split('', 1)[1].strip()
return line
return first_line if first_line else "未命名笔记"
def clean_content(self, content):
"""清理内容格式"""
if not content:
return ""
content = content.strip()
if not content:
return ""
# 移除文件名
content = re.sub(r'captured_text_\d{8}_\d{6}\n*', '', content)
# 使用结巴分词提取关键词
keywords = jieba.analyse.textrank(content[:500], topK=3, allowPOS=('ns', 'n', 'vn', 'v'))
if keywords:
return ' '.join(keywords)
# 清理多余的换行符,但保留段落格式
content = re.sub(r'\n{3,}', '\n\n', content)
# 如果无法提取关键词使用内容的前30个字符
return content[:30].strip() + "..."
# 确保段落之间有一个空行
paragraphs = [p.strip() for p in content.split('\n\n')]
content = '\n\n'.join(p for p in paragraphs if p)
return content.strip()
def extract_tags(self, content):
"""提取内容的标签"""
@ -101,6 +143,9 @@ class DocumentProcessor:
# 提取标题
title = self.extract_title(content)
# 清理内容
content = self.clean_content(content)
# 提取标签
tags = self.extract_tags(content)

View File

@ -1,5 +1,37 @@
"""测试数据模块"""
# 带文件名和多余换行符的内容
CONTENT_WITH_FILENAME = """captured_text_20250115_220948
这篇文章是一篇关于自然语言处理NLP和人工智能AI技术演进的讨论
文章概要
讨论了从传统的RAG技术到Graph RAG的技术演进
基于RAG的三个基本阶段索引检索生成抽象出通用的RAG框架支持多种索引形式向量全文
探讨了Graph RAG未来的优化方向包括内容索引和检索生成阶段的改进思路
提到了RAG技术可能向智能体Agent架构演进的趋势"""
# 带数字编号的内容
NUMBERED_CONTENT = """1. 引言
2. 技术背景
3. 主要内容
自然语言处理是人工智能的重要分支
4. 结论"""
# 长标题内容
LONG_TITLE_CONTENT = """这是一个非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常非常长的标题,这是第二句话。
这是正文内容"""
# 普通文本内容
PLAIN_TEXT = """这是一个测试文本
它包含多行内容

View File

@ -0,0 +1,64 @@
"""文档处理器测试模块"""
import unittest
from ..document_processor import DocumentProcessor
from . import test_data
class TestDocumentProcessor(unittest.TestCase):
"""文档处理器测试类"""
def setUp(self):
"""测试准备"""
self.processor = DocumentProcessor()
def test_extract_title_with_filename(self):
"""测试从带文件名的内容中提取标题"""
title = self.processor.extract_title(test_data.CONTENT_WITH_FILENAME)
self.assertEqual(
title,
"这篇文章是一篇关于自然语言处理NLP和人工智能AI技术演进的讨论"
)
def test_extract_title_from_numbered_content(self):
"""测试从带数字编号的内容中提取标题"""
title = self.processor.extract_title(test_data.NUMBERED_CONTENT)
self.assertEqual(title, "自然语言处理是人工智能的重要分支")
def test_extract_title_from_long_content(self):
"""测试从长标题内容中提取标题"""
title = self.processor.extract_title(test_data.LONG_TITLE_CONTENT)
print(f"\n实际标题: '{title}'") # 打印实际标题
print(f"标题长度: {len(title)}") # 打印标题长度
self.assertTrue(len(title) <= 50)
# 检查是否包含原始内容的开头部分
self.assertTrue(title.startswith("这是一个非常"))
# 检查是否正确截断
self.assertTrue(title.endswith('...'))
def test_clean_content_with_filename(self):
"""测试清理带文件名的内容"""
content = self.processor.clean_content(test_data.CONTENT_WITH_FILENAME)
self.assertFalse('captured_text_' in content)
self.assertEqual(content.count('\n\n'), 5) # 检查段落分隔
def test_clean_content_with_extra_newlines(self):
"""测试清理带多余换行符的内容"""
content = "第一段\n\n\n\n第二段\n\n\n\n\n第三段"
cleaned = self.processor.clean_content(content)
self.assertEqual(cleaned, "第一段\n\n第二段\n\n第三段")
def test_empty_content(self):
"""测试空内容处理"""
title = self.processor.extract_title("")
self.assertEqual(title, "未命名笔记")
content = self.processor.clean_content("")
self.assertEqual(content, "")
def test_whitespace_content(self):
"""测试空白内容处理"""
title = self.processor.extract_title(" \n \t \n ")
self.assertEqual(title, "未命名笔记")
content = self.processor.clean_content(" \n \t \n ")
self.assertEqual(content, "")
if __name__ == '__main__':
unittest.main()