改进匹配算法
This commit is contained in:
@@ -15,6 +15,11 @@ from datetime import datetime
|
||||
import time
|
||||
import re
|
||||
import shutil
|
||||
import requests
|
||||
import math
|
||||
|
||||
OLLAMA_EMBED_URL = os.environ.get('OLLAMA_EMBED_URL', 'http://localhost:11434/api/embeddings')
|
||||
OLLAMA_EMBED_MODEL = os.environ.get('OLLAMA_EMBED_MODEL', 'nomic/text-embedding-3-large')
|
||||
|
||||
|
||||
def clear_directory(dir_path: Path) -> None:
|
||||
@@ -27,6 +32,93 @@ def clear_directory(dir_path: Path) -> None:
|
||||
shutil.rmtree(item)
|
||||
|
||||
|
||||
def get_embedding(text: str) -> List[float]:
|
||||
"""使用 Ollama 获取文本的 embedding"""
|
||||
try:
|
||||
response = requests.post(
|
||||
OLLAMA_EMBED_URL,
|
||||
json={
|
||||
'model': OLLAMA_EMBED_MODEL,
|
||||
'input': text,
|
||||
},
|
||||
timeout=10,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if isinstance(data, dict):
|
||||
if 'embedding' in data:
|
||||
return data['embedding']
|
||||
if 'data' in data and isinstance(data['data'], list) and data['data']:
|
||||
return data['data'][0].get('embedding', []) or []
|
||||
return []
|
||||
except Exception as e:
|
||||
print(f"获取 embedding 失败: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
|
||||
"""计算两个向量的余弦相似度"""
|
||||
if not vec1 or not vec2:
|
||||
return 0.0
|
||||
dot = sum(a * b for a, b in zip(vec1, vec2))
|
||||
norm1 = math.sqrt(sum(a * a for a in vec1))
|
||||
norm2 = math.sqrt(sum(b * b for b in vec2))
|
||||
return dot / (norm1 * norm2) if norm1 and norm2 else 0.0
|
||||
|
||||
|
||||
def normalize_for_match(text: str) -> str:
|
||||
"""规范化文本,去掉路径和扩展名后的匹配用字符串"""
|
||||
text = str(text).lower()
|
||||
text = re.sub(r'[\u0020-\u002f\u003a-\u0040\u005b-\u0060\u007b-\u007e]+', ' ', text)
|
||||
text = re.sub(r'\s+', ' ', text).strip()
|
||||
return text
|
||||
|
||||
|
||||
def token_overlap_score(a: str, b: str) -> float:
|
||||
tokens_a = set(re.findall(r'[\u4e00-\u9fff]+|\d+|[a-z]+', a))
|
||||
tokens_b = set(re.findall(r'[\u4e00-\u9fff]+|\d+|[a-z]+', b))
|
||||
if not tokens_a or not tokens_b:
|
||||
return 0.0
|
||||
return len(tokens_a & tokens_b) / max(1, len(tokens_b))
|
||||
|
||||
|
||||
def find_best_match(title: str, candidates: List[Tuple[Path, datetime]]) -> Tuple[Path, datetime]:
|
||||
"""从候选文件中找到与 title 最相似的文件"""
|
||||
if not candidates:
|
||||
return None
|
||||
|
||||
title_norm = normalize_for_match(title)
|
||||
exact_candidates = []
|
||||
for file_path, ctime in candidates:
|
||||
stem_norm = normalize_for_match(file_path.stem)
|
||||
if stem_norm and (stem_norm == title_norm or stem_norm in title_norm or title_norm in stem_norm):
|
||||
exact_candidates.append((file_path, ctime, stem_norm))
|
||||
|
||||
if exact_candidates:
|
||||
# 直接返回最精确的“包含匹配”候选项
|
||||
return max(exact_candidates, key=lambda item: len(item[2]))[:2]
|
||||
|
||||
title_emb = get_embedding(title)
|
||||
best_match = None
|
||||
best_score = -1.0
|
||||
|
||||
for file_path, ctime in candidates:
|
||||
stem_norm = normalize_for_match(file_path.stem)
|
||||
overlap = token_overlap_score(title_norm, stem_norm)
|
||||
score = overlap * 0.3
|
||||
|
||||
if title_emb:
|
||||
stem_emb = get_embedding(file_path.stem)
|
||||
if stem_emb:
|
||||
score += cosine_similarity(title_emb, stem_emb) * 0.7
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_match = (file_path, ctime)
|
||||
|
||||
return best_match or candidates[0]
|
||||
|
||||
|
||||
def extract_all_zips(
|
||||
source_dir: Union[str, Path] = "cc",
|
||||
output_dir: Union[str, Path] = "data",
|
||||
@@ -365,6 +457,8 @@ def build_file_associations(
|
||||
# 提取两组数据:标题、网址、时间
|
||||
title1, url1, time1 = parts[0].strip(), parts[1].strip(), parts[2].strip()
|
||||
title2, url2, time2 = parts[3].strip(), parts[4].strip(), parts[5].strip()
|
||||
if title1=="2.1 楞次定律 课件-2022-2023学年高一下学期物理人教版(2019)选择性必修第二册" or title2=="2.1 楞次定律 课件-2022-2023学年高一下学期物理人教版(2019)选择性必修第二册":
|
||||
print(f"调试:标题1={title1}, 标题2={title2}, url1={url1}, url2={url2}, time1={time1}, time2={time2}")
|
||||
|
||||
# 辅助函数:处理文件路径,如果是 zip 则递归解压
|
||||
def process_file(file_path: Path) -> List[str]:
|
||||
@@ -395,7 +489,11 @@ def build_file_associations(
|
||||
|
||||
# 第一组:根据网址匹配文件
|
||||
if 'www.21cnjy.com' in url1 and twole_files:
|
||||
file_path, ctime = twole_files.pop(0)
|
||||
candidates = twole_files
|
||||
best_match = find_best_match(title1, candidates)
|
||||
if best_match:
|
||||
twole_files.remove(best_match)
|
||||
file_path, ctime = best_match
|
||||
result.append({
|
||||
"filename": process_file(file_path),
|
||||
"title": title1,
|
||||
@@ -403,7 +501,11 @@ def build_file_associations(
|
||||
"time": time1
|
||||
})
|
||||
elif 'www.zxxk.com' in url1 and zxxk_files:
|
||||
file_path, ctime = zxxk_files.pop(0)
|
||||
candidates = zxxk_files
|
||||
best_match = find_best_match(title1, candidates)
|
||||
if best_match:
|
||||
zxxk_files.remove(best_match)
|
||||
file_path, ctime = best_match
|
||||
result.append({
|
||||
"filename": process_file(file_path),
|
||||
"title": title1,
|
||||
@@ -413,7 +515,11 @@ def build_file_associations(
|
||||
|
||||
# 第二组:根据网址匹配文件
|
||||
if 'www.21cnjy.com' in url2 and twole_files:
|
||||
file_path, ctime = twole_files.pop(0)
|
||||
candidates = twole_files
|
||||
best_match = find_best_match(title2, candidates)
|
||||
if best_match:
|
||||
twole_files.remove(best_match)
|
||||
file_path, ctime = best_match
|
||||
result.append({
|
||||
"filename": process_file(file_path),
|
||||
"title": title2,
|
||||
@@ -421,7 +527,11 @@ def build_file_associations(
|
||||
"time": time2
|
||||
})
|
||||
elif 'www.zxxk.com' in url2 and zxxk_files:
|
||||
file_path, ctime = zxxk_files.pop(0)
|
||||
candidates = zxxk_files
|
||||
best_match = find_best_match(title2, candidates)
|
||||
if best_match:
|
||||
zxxk_files.remove(best_match)
|
||||
file_path, ctime = best_match
|
||||
result.append({
|
||||
"filename": process_file(file_path),
|
||||
"title": title2,
|
||||
|
||||
Reference in New Issue
Block a user