94 lines
1.8 KiB
Python
94 lines
1.8 KiB
Python
"""稀疏检索 / 简单 RAG"""
|
||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
from sklearn.metrics.pairwise import cosine_similarity
|
||
|
||
|
||
# 1. 简单切块
|
||
def chunk_text(text, chunk_size=60, overlap=15):
|
||
chunks = []
|
||
|
||
start = 0
|
||
|
||
while start < len(text):
|
||
end = start + chunk_size
|
||
|
||
chunk = text[start:end]
|
||
chunks.append(chunk)
|
||
|
||
if end >= len(text):
|
||
break
|
||
|
||
start = end - overlap
|
||
|
||
return chunks
|
||
|
||
|
||
# 2. 知识库原文
|
||
document = """
|
||
职工请事假，应当提前一个工作日通过OA系统提交申请，
|
||
经部门负责人审批同意后方可休假。
|
||
|
||
职工因疾病无法正常工作的，可以申请病假。
|
||
病假应当提供医院出具的相关证明材料。
|
||
|
||
员工因工作需要出差的，应提前提交出差申请。
|
||
住宿费用按照公司差旅管理制度规定的标准报销。
|
||
|
||
职工连续工作满一年后可以享受带薪年休假。
|
||
年休假天数根据职工累计工作年限确定。
|
||
""".replace("\n", "")
|
||
|
||
|
||
# 3. 文档切块
|
||
chunks = chunk_text(
|
||
document,
|
||
chunk_size=60,
|
||
overlap=15
|
||
)
|
||
|
||
print("=== chunks ===")
|
||
|
||
for i, chunk in enumerate(chunks):
|
||
print(i, chunk)
|
||
|
||
|
||
# 4. 给所有 chunk 建 TF-IDF 索引
|
||
vectorizer = TfidfVectorizer(
|
||
analyzer="char",
|
||
ngram_range=(1, 3)
|
||
)
|
||
|
||
doc_vectors = vectorizer.fit_transform(chunks)
|
||
|
||
|
||
# 5. 用户问题
|
||
query = "事假怎么申请？"
|
||
|
||
|
||
# 6. Query 转成同一个 TF-IDF 空间
|
||
query_vector = vectorizer.transform([query])
|
||
|
||
|
||
# 7. 计算 query 与所有 chunk 的相似度
|
||
scores = cosine_similarity(
|
||
query_vector,
|
||
doc_vectors
|
||
)[0]
|
||
|
||
|
||
# 8. Top-K
|
||
top_k = 3
|
||
|
||
top_indices = scores.argsort()[::-1][:top_k]
|
||
|
||
|
||
print("\n=== query ===")
|
||
print(query)
|
||
|
||
print("\n=== retrieval result ===")
|
||
|
||
for index in top_indices:
|
||
print(
|
||
f"score={scores[index]:.4f}",
|
||
chunks[index]
|
||
) |