knowledge/utils/text_similarity.py

import os
import numpy as np
import re
from typing import List, Dict, Tuple, Optional, Union, Callable
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import jieba
import argparse

class TextSimilarityFinder:
    """文本相似度查找器，用于从多个文本中找到与目标文本最相似的文本"""
    
    def __init__(self, method: str = 'tfidf', use_jieba: bool = True, ngram_range: Tuple[int, int] = (1, 1)):
        """
        初始化文本相似度查找器
        
        Args:
            method: 相似度计算方法，可选值为 'tfidf'、'count'、'jaccard'
            use_jieba: 是否使用结巴分词（中文文本建议开启）
            ngram_range: N-gram范围，例如(1,2)表示同时使用unigram和bigram特征
        """
        self.method = method
        self.use_jieba = use_jieba
        self.ngram_range = ngram_range
        self.vectorizer = None
        self.corpus_vectors = None
        self.corpus_texts = None
        self.corpus_paths = None
    
    def _tokenize_text(self, text: str) -> str:
        """对文本进行分词处理
        
        Args:
            text: 输入文本
            
        Returns:
            str: 分词后的文本
        """
        if not self.use_jieba:
            return text
        
        # 使用结巴分词，返回空格分隔的词语
        return " ".join(jieba.cut(text))
    
    def _preprocess_texts(self, texts: List[str]) -> List[str]:
        """预处理文本列表
        
        Args:
            texts: 文本列表
            
        Returns:
            List[str]: 预处理后的文本列表
        """
        processed_texts = []
        for text in texts:
            # 去除标点符号、特殊字符和多余空白字符
            # 匹配中文标点符号（包括圆点•）
            text = re.sub(r'[\u3000-\u303F]|[\uFF00-\uFFEF]|[•]', ' ', text)
            # 匹配英文标点符号和特殊字符
            text = re.sub(r'[!"#$%&\'()*+,-./:;<=>?@\[\\\]^_`{|}~]', ' ', text)
            # 匹配其他特殊字符（如制表符、换行符等）
            text = re.sub(r'[\t\n\r\f\v]', ' ', text)
            # 去除多余空白字符
            text = re.sub(r'\s+', ' ', text).strip()
            # 分词处理
            text = self._tokenize_text(text)
            processed_texts.append(text)
        return processed_texts
    
    def _initialize_vectorizer(self):
        """初始化文本向量化器"""
        if self.method == 'tfidf':
            self.vectorizer = TfidfVectorizer(ngram_range=self.ngram_range)
        elif self.method == 'count':
            self.vectorizer = CountVectorizer(ngram_range=self.ngram_range)
        elif self.method == 'jaccard':
            # Jaccard相似度不需要向量化器，使用集合操作
            self.vectorizer = None
        else:
            raise ValueError(f"不支持的相似度计算方法: {self.method}")
    
    def _calculate_jaccard_similarity(self, text1: str, text2: str) -> float:
        """计算两个文本的Jaccard相似度
        
        Args:
            text1: 第一个文本
            text2: 第二个文本
            
        Returns:
            float: Jaccard相似度
        """
        # 将文本转换为词集合
        set1 = set(text1.split())
        set2 = set(text2.split())
        
        # 计算Jaccard相似度: 交集大小 / 并集大小
        intersection = len(set1.intersection(set2))
        union = len(set1.union(set2))
        
        # 避免除以零
        if union == 0:
            return 0.0
        
        return intersection / union
    
    def load_corpus_from_directory(self, directory: str, file_pattern: str = '*.txt'):
        """从目录加载语料库
        
        Args:
            directory: 目录路径
            file_pattern: 文件匹配模式
        """
        import glob
        
        # 获取所有匹配的文件路径
        file_paths = glob.glob(os.path.join(directory, file_pattern))
        
        if not file_paths:
            raise ValueError(f"在目录 {directory} 中没有找到匹配 {file_pattern} 的文件")
        
        # 读取所有文件内容
        texts = []
        for file_path in file_paths:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    texts.append(f.read())
            except Exception as e:
                print(f"读取文件 {file_path} 失败: {str(e)}")
        
        self.load_corpus(texts, file_paths)
    
    def load_corpus(self, texts: List[str], paths: Optional[List[str]] = None):
        """加载语料库
        
        Args:
            texts: 文本列表
            paths: 文本对应的路径或标识符列表
        """
        if not texts:
            raise ValueError("文本列表不能为空")
        
        # 预处理文本
        processed_texts = self._preprocess_texts(texts)
        
        # 初始化向量化器
        self._initialize_vectorizer()
        
        # 存储原始文本和路径
        self.corpus_texts = texts
        self.corpus_paths = paths if paths else [f"text_{i}" for i in range(len(texts))]
        
        # 向量化语料库
        if self.method in ['tfidf', 'count']:
            self.corpus_vectors = self.vectorizer.fit_transform(processed_texts)
        else:
            # 对于Jaccard相似度，直接存储处理后的文本
            self.corpus_vectors = processed_texts
    
    def find_most_similar(self, query_text: str, top_n: int = 1) -> List[Dict]:
        """查找与查询文本最相似的文本
        
        Args:
            query_text: 查询文本
            top_n: 返回最相似的前N个结果
            
        Returns:
            List[Dict]: 包含相似度信息的结果列表
        """
        if self.corpus_vectors is None:
            raise ValueError("请先加载语料库")
        
        # 预处理查询文本
        processed_query = self._preprocess_texts([query_text])[0]
        
        # 计算相似度
        similarities = []
        
        if self.method in ['tfidf', 'count']:
            # 向量化查询文本
            query_vector = self.vectorizer.transform([processed_query])
            
            # 计算余弦相似度
            similarity_matrix = cosine_similarity(query_vector, self.corpus_vectors)
            similarities = similarity_matrix[0]
        else:  # Jaccard相似度
            for corpus_text in self.corpus_vectors:
                similarity = self._calculate_jaccard_similarity(processed_query, corpus_text)
                similarities.append(similarity)
        
        # 获取相似度排序的索引
        top_indices = np.argsort(similarities)[::-1][:top_n]
        
        # 构建结果
        results = []
        for idx in top_indices:
            results.append({
                'text': self.corpus_texts[idx],
                'path': self.corpus_paths[idx],
                'similarity': similarities[idx]
            })
        
        return results

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='文本相似度查找工具')
    parser.add_argument('--query', type=str, required=True, help='查询文本或查询文本文件路径')
    parser.add_argument('--dir', type=str, required=True, help='语料库目录路径')
    parser.add_argument('--pattern', type=str, default='*.txt', help='文件匹配模式')
    parser.add_argument('--method', type=str, default='tfidf', choices=['tfidf', 'count', 'jaccard'], help='相似度计算方法')
    parser.add_argument('--top_n', type=int, default=3, help='返回最相似的前N个结果')
    parser.add_argument('--no_jieba', action='store_true', help='不使用结巴分词')
    parser.add_argument('--ngram_min', type=int, default=1, help='N-gram最小值')
    parser.add_argument('--ngram_max', type=int, default=1, help='N-gram最大值')
    
    args = parser.parse_args()
    
    # 初始化文本相似度查找器
    finder = TextSimilarityFinder(
        method=args.method,
        use_jieba=not args.no_jieba,
        ngram_range=(args.ngram_min, args.ngram_max)
    )
    
    # 加载语料库
    finder.load_corpus_from_directory(args.dir, args.pattern)
    
    # 获取查询文本
    query_text = args.query
    if os.path.isfile(query_text):
        try:
            with open(query_text, 'r', encoding='utf-8') as f:
                query_text = f.read()
        except Exception as e:
            print(f"读取查询文件失败: {str(e)}")
            return
    
    # 查找最相似的文本
    results = finder.find_most_similar(query_text, args.top_n)
    
    # 打印结果
    print(f"\n查询文本: {query_text[:100]}..." if len(query_text) > 100 else f"\n查询文本: {query_text}")
    print(f"\n找到 {len(results)} 个最相似的文本:")
    
    for i, result in enumerate(results):
        print(f"\n{i+1}. 相似度: {result['similarity']:.4f}")
        print(f"   路径: {result['path']}")
        text_preview = result['text'][:200] + '...' if len(result['text']) > 200 else result['text']
        print(f"   内容预览: {text_preview}")

if __name__ == "__main__":
    main()