"""第 32 章：Markdown 片段、BM25、实测稠密向量与 RRF；默认不下载模型。""" from collections import Counter import math import re from prompt_contracts import read_json, fingerprint EMBEDDING_MODEL = "BAAI/bge-small-zh-v1.5" EMBEDDING_REVISION = "7999e1d3359715c523056ef9478215996d62a620" QUERY_PREFIX = "为这个句子生成表示以用于检索相关文章：" def chunk_documents(documents, max_chars=180): """只处理本书使用的二级标题与纯文本段落，保留正文字符偏移。""" if max_chars < 1: raise ValueError("块长必须为正") chunks = [] for doc in documents: # 本例只有公开读者这一种身份；生产系统需由服务端身份决定 ACL。 if not doc["current"] or doc["visibility"] != "public": continue body = doc["body"] sections = list(re.finditer(r"(?m)^## ([^\n]+)\n([\s\S]*?)(?=^## |\Z)", body)) if not sections or body[:sections[0].start()].strip(): raise ValueError("本例要求正文从二级标题开始") part = 0 for section in sections: for paragraph in re.finditer(r"[^\n]+", section.group(2)): start = section.start(2)+paragraph.start() remaining = paragraph.group() while remaining: end = min(max_chars, len(remaining)) if end < len(remaining): boundary = remaining.rfind("。", 0, end) if boundary >= 0: end = boundary+1 content = remaining[:end] part += 1 chunks.append(dict( id=f"{doc['id']}:{part}", document=doc["id"], version=doc["version"], start=start, end=start+end, content=content, text=f"{doc['title']} / {section.group(1)}\n{content}", document_sha256=fingerprint(body), )) start += end remaining = remaining[end:] if len({chunk["id"] for chunk in chunks}) != len(chunks): raise ValueError("片段 ID 重复") return chunks def terms(text): result = [] for run in re.findall(r"[\u4e00-\u9fff]+|[a-z0-9]+", text.lower()): if re.fullmatch(r"[a-z0-9]+", run): result.append(run) else: result.extend(run) result.extend(run[i:i+2] for i in range(len(run)-1)) return result def bm25(query, chunks, k1=1.2, b=.75): counts = [Counter(terms(chunk["text"])) for chunk in chunks] lengths = [sum(count.values()) for count in counts] average = sum(lengths)/len(lengths) if lengths else 1 scores = [0.0]*len(chunks) for term in set(terms(query)): df = sum(term in count for count in counts) idf = math.log(1+(len(chunks)-df+.5)/(df+.5)) for i, (count, length) in enumerate(zip(counts, lengths)): tf = count[term] denominator = tf+k1*(1-b+b*length/average) scores[i] += idf*tf*(k1+1)/denominator if denominator else 0 return sorted(zip(chunks, scores), key=lambda row: (-row[1], row[0]["id"])) def vector_inputs(chunks, questions): return ([dict(key=chunk["id"], text=chunk["text"]) for chunk in chunks] + [dict(key=case["id"], text=QUERY_PREFIX+case["question"]) for case in questions]) def load_vectors(chunks, questions, name="workshop-vectors.json"): saved = read_json(name) expected = vector_inputs(chunks, questions) if (saved["inputs_sha256"] != fingerprint(expected) or saved["model"] != EMBEDDING_MODEL or saved["revision"] != EMBEDDING_REVISION): raise ValueError("向量对应的文本或模型已变化，请重新生成向量") vectors = saved["vectors"] if set(vectors) != {item["key"] for item in expected}: raise ValueError("向量条目不完整") for vector in vectors.values(): if len(vector) != 512 or any(not math.isfinite(x) for x in vector): raise ValueError("向量形状或数值错误") if abs(sum(x*x for x in vector)-1) > 1e-5: raise ValueError("向量未归一化") return vectors def dense(case, chunks, vectors): q = vectors[case["id"]] return sorted(((chunk, sum(a*b for a, b in zip(q, vectors[chunk["id"]]))) for chunk in chunks), key=lambda row: (-row[1], row[0]["id"])) def reciprocal_rank_fusion(rankings, c=60, candidate_k=5): scores, by_id = Counter(), {} for ranking in rankings: for position, (chunk, _) in enumerate(ranking[:candidate_k], 1): scores[chunk["id"]] += 1/(c+position) by_id[chunk["id"]] = chunk return [(by_id[key], score) for key, score in sorted(scores.items(), key=lambda row: (-row[1], row[0]))] def retrieve(case, chunks, vectors, method="rrf", k=2): sparse = bm25(case["question"], chunks) neural = dense(case, chunks, vectors) ranked = {"bm25": sparse, "dense": neural, "rrf": reciprocal_rank_fusion([sparse, neural])} if method not in ranked: raise ValueError("未知检索方法") return [chunk for chunk, _ in ranked[method][:k]] def retrieval_metrics(cases, rankings): answerable = [case for case in cases if case["gold"]] if not answerable: raise ValueError("没有有答案的问题，召回率无定义") recall = complete = hit1 = 0 for case in answerable: expected = set(case["gold"]) found = [chunk["id"] for chunk in rankings[case["id"]]] recall += len(expected.intersection(found))/len(expected) complete += expected.issubset(found) hit1 += bool(found and found[0] in expected) count = len(answerable) return dict(count=count, recall=recall/count, complete=complete, hit1=hit1) def run(): docs = read_json("workshop-documents.json") chunks = chunk_documents(docs) questions = read_json("workshop-questions.json") vectors = load_vectors(chunks, questions) cases = [case for case in questions if case["split"] == "test"] report = {} for method in ["bm25", "dense", "rrf"]: rankings = {case["id"]: retrieve(case, chunks, vectors, method) for case in cases} report[method] = dict(metrics=retrieval_metrics(cases, rankings), rankings={key: [row["id"] for row in value] for key, value in rankings.items()}) return dict(documents=len(docs), chunks=len(chunks), methods=report) def main(): report = run() print(f"原始文档={report['documents']}；公开且生效的片段={report['chunks']}；每次取 2 块") for method, row in report["methods"].items(): m = row["metrics"] print(f"{method}：首位命中={m['hit1']}/{m['count']}；" f"平均 Recall@2={m['recall']:.3f}；证据齐全={m['complete']}/{m['count']}") print("多来源问题："+", ".join(report["methods"]["rrf"]["rankings"]["test-two"])) print("无答案问题仍会返回："+", ".join(report["methods"]["rrf"]["rankings"]["test-unknown"])) print("排名靠前不等于资料中有答案；引用存在不等于引用支持结论。") if __name__ == "__main__": main()