116 lines
3.3 KiB
Python
116 lines
3.3 KiB
Python
"""
|
||
Blog Semantic Search
|
||
====================
|
||
讀取 blog_embeddings.json，讓使用者輸入（或貼上）任意文字，
|
||
用同一個 embedding 模型算出向量，列出最相關的 20 篇文章。
|
||
|
||
需求：
|
||
- ollama 已啟動且已 pull qwen3-embedding:8b
|
||
- pip install numpy requests
|
||
- 已經跑過 blog_embeddings.py 產生 blog_embeddings.json
|
||
"""
|
||
|
||
import json
|
||
import numpy as np
|
||
import requests
|
||
|
||
# ============================================================
|
||
# 設定區
|
||
# ============================================================
|
||
|
||
EMBEDDINGS_FILE = "./blog_embeddings.json"
|
||
OLLAMA_URL = "http://localhost:11434/api/embed"
|
||
OLLAMA_MODEL = "qwen3-embedding:8b"
|
||
|
||
# ============================================================
|
||
|
||
|
||
def get_embedding(text: str) -> list[float]:
|
||
"""透過 Ollama API 取得 embedding 向量。"""
|
||
resp = requests.post(OLLAMA_URL, json={
|
||
"model": OLLAMA_MODEL,
|
||
"input": text,
|
||
})
|
||
resp.raise_for_status()
|
||
return resp.json()["embeddings"][0]
|
||
|
||
|
||
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
|
||
"""計算兩個向量的 cosine similarity。"""
|
||
dot = np.dot(a, b)
|
||
norm = np.linalg.norm(a) * np.linalg.norm(b)
|
||
if norm == 0:
|
||
return 0.0
|
||
return float(dot / norm)
|
||
|
||
|
||
def format_row(rank: int, sim: float, title: str, slug: str) -> str:
|
||
"""格式化一行結果。"""
|
||
bar = "█" * int(sim * 30)
|
||
return f" {rank:3d}. {sim:.4f} {bar} {title} ({slug})"
|
||
|
||
|
||
def read_query() -> str:
|
||
"""讀取使用者輸入，支援多行貼上（連按兩次 Enter 結束）。"""
|
||
print("請輸入搜尋文字（連按兩次 Enter 結束）：")
|
||
lines = []
|
||
empty_count = 0
|
||
while True:
|
||
line = input()
|
||
if line == "":
|
||
empty_count += 1
|
||
if empty_count >= 2:
|
||
break
|
||
lines.append(line)
|
||
else:
|
||
empty_count = 0
|
||
lines.append(line)
|
||
return "\n".join(lines).strip()
|
||
|
||
|
||
def main():
|
||
# 讀取 embeddings
|
||
with open(EMBEDDINGS_FILE, encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
if not data:
|
||
print("找不到任何文章資料。")
|
||
return
|
||
|
||
print(f"已載入 {len(data)} 篇文章的 embeddings\n")
|
||
|
||
# 預先轉成 numpy array，之後比較比較快
|
||
vectors = [np.array(item["embedding"]) for item in data]
|
||
|
||
while True:
|
||
query = read_query()
|
||
if not query:
|
||
print("沒有輸入，掰掰！")
|
||
break
|
||
|
||
print(f"\n正在產生 embedding（模型：{OLLAMA_MODEL}）...")
|
||
query_vec = np.array(get_embedding(query))
|
||
|
||
# 計算與所有文章的相似度
|
||
similarities = []
|
||
for i, item in enumerate(data):
|
||
sim = cosine_similarity(query_vec, vectors[i])
|
||
similarities.append({
|
||
"slug": item["slug"],
|
||
"title": item.get("title", "(無標題)"),
|
||
"similarity": sim,
|
||
})
|
||
|
||
similarities.sort(key=lambda x: x["similarity"], reverse=True)
|
||
|
||
top = similarities[:20]
|
||
print(f"\n🔍 與輸入文字最相關的 20 篇：\n")
|
||
for rank, s in enumerate(top, 1):
|
||
print(format_row(rank, s["similarity"], s["title"], s["slug"]))
|
||
|
||
print("\n" + "─" * 60 + "\n")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|