import json import re import faiss import numpy as np from pathlib import Path from typing import List, Dict, Set from .embedding import EmbeddingEngine class PoetryMatcher: def __init__(self, data_path: str = "data/quotes.json", index_path: str = "data/index.faiss", synonym_path: str = "data/synonym_dict.json", phrase_path: str = "data/phrase_dict.json", model_name: str = "moka-ai/m3e-base", semantic_weight: float = 0.5, keyword_weight: float = 0.25, synonym_weight: float = 0.25): self.data_path = Path(data_path) self.index_path = Path(index_path) self.synonym_path = Path(synonym_path) self.phrase_path = Path(phrase_path) self.semantic_weight = semantic_weight self.keyword_weight = keyword_weight self.synonym_weight = synonym_weight self.embedding_engine = EmbeddingEngine(model_name) self.quotes = self._load_quotes_from_json() print(f"Loaded {len(self.quotes)} quotes from JSON") self.synonym_dict = self._load_synonym_dict() print(f"Loaded synonym dict with {len(self.synonym_dict)} entries") self.phrase_dict = self._load_phrase_dict() print(f"Loaded phrase dict with {len(self.phrase_dict)} entries") self.index = self._load_or_create_index() def _load_quotes_from_json(self) -> List[Dict]: if not self.data_path.exists(): print(f"Data file not found: {self.data_path}") return [] with open(self.data_path, 'r', encoding='utf-8') as f: quotes = json.load(f) return quotes def _load_synonym_dict(self) -> Dict[str, List[str]]: """加载同义词词典""" if not self.synonym_path.exists(): print(f"Synonym dict not found: {self.synonym_path}") return {} with open(self.synonym_path, 'r', encoding='utf-8') as f: raw_dict = json.load(f) # 将分类词典转换为扁平化的映射 synonym_map = {} for category, mappings in raw_dict.items(): for word, synonyms in mappings.items(): synonym_map[word] = synonyms return synonym_map def _load_phrase_dict(self) -> Dict[str, List[str]]: """加载短语词典""" if not self.phrase_path.exists(): print(f"Phrase dict not found: {self.phrase_path}") return {} with open(self.phrase_path, 'r', encoding='utf-8') as f: raw_dict = json.load(f) # 将分类词典转换为扁平化的映射 phrase_map = {} for category, mappings in raw_dict.items(): for phrase, synonyms in mappings.items(): phrase_map[phrase] = synonyms return phrase_map def _expand_query(self, query: str) -> List[str]: """扩展查询词""" expanded = [query] # 原始查询 # 1. 优先检查短语匹配（从长到短） sorted_phrases = sorted(self.phrase_dict.keys(), key=len, reverse=True) for phrase in sorted_phrases: if phrase in query: synonyms = self.phrase_dict[phrase] expanded.extend(synonyms) # 2. 提取中文词，检查单字同义词 words = re.findall(r'[\u4e00-\u9fff]+', query) for word in words: # 检查是否有同义词映射 if word in self.synonym_dict: synonyms = self.synonym_dict[word] expanded.extend(synonyms) # 检查2-gram if len(word) >= 2: for i in range(len(word) - 1): bigram = word[i:i+2] if bigram in self.synonym_dict: synonyms = self.synonym_dict[bigram] expanded.extend(synonyms) # 去重 return list(set(expanded)) def _extract_keywords(self, text: str) -> Set[str]: """从文本中提取关键词（单字+2-gram）""" chars = re.findall(r'[\u4e00-\u9fff]', text) keywords = set() for char in chars: keywords.add(char) for i in range(len(chars) - 1): bigram = chars[i] + chars[i+1] keywords.add(bigram) return keywords def _calculate_keyword_score(self, query: str, quote_text: str) -> float: """计算关键词匹配分数""" query_keywords = self._extract_keywords(query) quote_keywords = self._extract_keywords(quote_text) if not query_keywords: return 0.0 intersection = query_keywords & quote_keywords score = 0.0 for kw in intersection: if len(kw) >= 2: score += 2.0 else: score += 1.0 max_score = sum(2.0 if len(kw) >= 2 else 1.0 for kw in query_keywords) normalized_score = score / max_score if max_score > 0 else 0.0 return min(1.0, normalized_score) def _calculate_synonym_score(self, expanded_queries: List[str], quote_text: str) -> float: """计算同义词匹配分数""" if len(expanded_queries) <= 1: return 0.0 # 提取名句中的字符 quote_chars = set(re.findall(r'[\u4e00-\u9fff]', quote_text)) # 计算扩展词在名句中的命中率 hit_count = 0 for eq in expanded_queries[1:]: # 跳过原始查询 # 检查扩展词是否在名句中出现 if eq in quote_text: hit_count += 1 # 检查扩展词的每个字符是否在名句中 elif any(char in quote_chars for char in eq): hit_count += 0.5 # 归一化 max_hits = len(expanded_queries) - 1 if max_hits == 0: return 0.0 return min(1.0, hit_count / max_hits) def _load_or_create_index(self) -> faiss.Index: if self.index_path.exists(): print(f"Loading existing index from {self.index_path}") return faiss.read_index(str(self.index_path)) else: print("No existing index found, building now...") self.build_index() return self.index def build_index(self): if len(self.quotes) == 0: print("No quotes to index") return print(f"Building index for {len(self.quotes)} quotes...") batch_size = 32 all_embeddings = [] for i in range(0, len(self.quotes), batch_size): batch = self.quotes[i:i+batch_size] texts = [q['text'] for q in batch] embeddings = self.embedding_engine.encode(texts, use_cache=False) all_embeddings.append(embeddings) if (i // batch_size) % 10 == 0: print(f"Processed {min(i+batch_size, len(self.quotes))}/{len(self.quotes)}") all_embeddings = np.vstack(all_embeddings).astype('float32') dimension = all_embeddings.shape[1] self.index = faiss.IndexFlatIP(dimension) self.index.add(all_embeddings) self.index_path.parent.mkdir(parents=True, exist_ok=True) faiss.write_index(self.index, str(self.index_path)) print(f"Index saved to {self.index_path}") print(f"Index contains {self.index.ntotal} vectors") def match(self, query: str, top_k: int = 5) -> List[Dict]: if self.index is None or self.index.ntotal == 0: return [] # 扩展查询 expanded_queries = self._expand_query(query) # 语义搜索 semantic_top_k = min(top_k * 5, 100) query_embedding = self.embedding_engine.encode_single(query) query_array = np.array([query_embedding], dtype=np.float32) semantic_scores, indices = self.index.search(query_array, semantic_top_k) results = [] for sem_score, idx in zip(semantic_scores[0], indices[0]): if idx >= 0 and idx < len(self.quotes): quote = self.quotes[idx].copy() semantic_score = float(sem_score) keyword_score = self._calculate_keyword_score(query, quote['text']) synonym_score = self._calculate_synonym_score(expanded_queries, quote['text']) # 综合分数 final_score = (self.semantic_weight * semantic_score + self.keyword_weight * keyword_score + self.synonym_weight * synonym_score) # 构建返回数据，确保所有字段都存在 result = { 'id': quote.get('id', f'generated_{idx}'), 'text': quote.get('text', ''), 'author': quote.get('author'), 'source': quote.get('source'), 'dynasty': quote.get('dynasty'), 'type': quote.get('type'), 'score': final_score, 'semantic_score': semantic_score, 'keyword_score': keyword_score, 'synonym_score': synonym_score } results.append(result) results.sort(key=lambda x: x['score'], reverse=True) return results[:top_k] def get_quotes_count(self) -> int: return len(self.quotes)