Python根据关联词搜索模型

这是一段关联词汇搜索引擎代码，能根据词语共现关系扩展查询，让搜索结果更贴合语义。

```python

-*- coding: utf-8 -*-

"""

基于关联词汇扩展的搜索

思路：

1. 建索引：分词 -> 倒排索引

2. 挖关联词：文档内滑动窗口共现 -> PMI 计算关联强度 -> 归一化取 topk

3. 查询扩展：原词权重 1.0 + 同义词 β + 关联词 α * 关联度

4. BM25 打分排序

"""

import re

import math

from collections import defaultdict, Counter

---------- 分词器 ----------

try:

import jieba

def default_tokenizer(text):

return w.strip().lower() for w in jieba.lcut(text) if w.strip()

except ImportError:

_SEG_RE = re.compile(r"a-zA-Z0-9+|\\u4e00-\\u9fa5+")

def default_tokenizer(text):

"""无 jieba 时的兜底：中文切单字 + 双字，英文按单词"""

tokens = \[\]

for seg in _SEG_RE.findall(text.lower()):

if seg.isascii():

tokens.append(seg)

else:

tokens.extend(seg) # 单字

tokens.extend(segi:i + 2 for i in range(len(seg) - 1)) # 双字

return tokens

class RelatedSearchEngine:

"""基于关联词扩展的搜索引擎"""

def init(self, tokenizer=None, k1=1.5, b=0.75):

self.tokenizer = tokenizer or default_tokenizer

self.k1 = k1 # BM25 词频饱和参数

self.b = b # BM25 长度归一化参数

self.docs = {} # doc_id -> 原文

self.doc_len = {} # doc_id -> 词数

self.inverted = defaultdict(dict) # term -> {doc_id: tf}

self.df = Counter() # term -> 文档频率

self.idf = {}

self.avg_len = 0.0

self.related = {} # term -> {关联词: 权重 0~1}

self.synonyms = defaultdict(set) # term -> 同义词集合

self._built = False

---------------- 索引 ----------------

def add_document(self, doc_id, text):

terms = self.tokenizer(text)

self.docsdoc_id = text

self.doc_lendoc_id = len(terms)

for t, c in Counter(terms).items():

self.invertedtdoc_id = c

self.dft += 1

self._built = False

def add_documents(self, docs):

for doc_id, text in docs.items():

self.add_document(doc_id, text)

def add_synonyms(self, word, *synonyms):

"""手工补充同义词，双向生效"""

self.synonymsword.update(synonyms)

for s in synonyms:

self.synonymss.add(word)

def build(self, window=5, topk=10, min_co=2):

"""构建 idf 与关联词表"""

n = len(self.docs)

if n == 0:

raise ValueError("请先添加文档")

self.avg_len = sum(self.doc_len.values()) / n or 1.0

self.idf = {

t: math.log(1 + (n - df + 0.5) / (df + 0.5))

for t, df in self.df.items()

}

self._build_related(window, topk, min_co)

self._built = True

return self

def _build_related(self, window, topk, min_co):

"""滑窗共现 + PMI 挖掘关联词"""

co = defaultdict(Counter)

for text in self.docs.values():

seq = self.tokenizer(text)

L = len(seq)

for i, w in enumerate(seq):

for j in range(i + 1, min(i + window + 1, L)):

v = seqj

if v == w:

continue

cowv += 1

covw += 1

freq = {t: sum(c.values()) for t, c in co.items()}

total_pairs = sum(freq.values()) / 2 or 1 # 对称计数，除以 2

self.related = {}

for w, cnt in co.items():

raw = {}

for v, c in cnt.items():

if c < min_co:

continue

PMI：共现概率 / 各自独立出现概率

rawv = math.log(c * total_pairs / (freqw * freqv) + 1e-12)

if not raw:

continue

lo, hi = min(raw.values()), max(raw.values())

span = hi - lo

norm = {v: (s - lo) / span if span > 1e-12 else 1.0

for v, s in raw.items()}

self.relatedw = dict(sorted(norm.items(), key=lambda x: -x1):topk)

---------------- 检索 ----------------

def expand(self, query, alpha=0.5, beta=0.8):

"""把查询词扩展成 {词: 权重}"""

weights = Counter()

for t in self.tokenizer(query):

weightst = max(weightst, 1.0) # 原始词

for s in self.synonyms.get(t, ()):

weightss = max(weightss, beta) # 同义词

for r, sim in self.related.get(t, {}).items():

weightsr = max(weightsr, alpha * sim) # 关联词

return weights

def search(self, query, topn=5, alpha=0.5, beta=0.8, explain=False):

if not self._built:

self.build()

weights = self.expand(query, alpha, beta)

scores = defaultdict(float)

hits = defaultdict(list)

for term, wq in weights.items():

postings = self.inverted.get(term)

if not postings:

continue

idf = self.idf.get(term, 0.0)

for doc_id, tf in postings.items():

dl = self.doc_lendoc_id

denom = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_len)

gain = wq * idf * tf * (self.k1 + 1) / denom

scoresdoc_id += gain

if explain:

hitsdoc_id.append((term, round(gain, 3)))

ranked = sorted(scores.items(), key=lambda x: -x1):topn

if explain:

return (d, round(s, 4), sorted(hits\[d, key=lambda x: -x1))

for d, s in ranked]

return (d, round(s, 4)) for d, s in ranked

if name == "main":

docs = {

"d1": "苹果发布了新一代 iPhone，搭载全新的芯片和摄像头",

"d2": "iPhone 的电池续航一直是用户关心的问题",

"d3": "华为发布了新款手机，影像系统升级明显",

"d4": "小米手机的性价比很高，适合学生党",

"d5": "苹果公司的股价今天上涨了百分之三",

"d6": "手机摄影技巧：如何拍出好看的夜景照片",

}

engine = RelatedSearchEngine()

engine.add_documents(docs)

engine.add_synonyms("手机", "电话", "移动电话")

engine.build(window=5, topk=8, min_co=1)

for q in "苹果", "手机拍照", "iPhone":

print(f"查询：{q}")

for doc_id, score, detail in engine.search(q, topn=3, explain=True):

print(f" {doc_id} {score} {docsdoc_id:22}...")

print(f" 命中：{detail}")

print()

查看自动挖出的关联词

print("『手机』的关联词：", engine.related.get("手机"))

print("『苹果』的关联词：", engine.related.get("苹果"))

```

关联搜索的核心模块与交互流程

代码把搜索拆成几个关键步骤，让关联扩展和打分清晰可控。

· 索引与关联词挖掘：先为每个文档分词建立倒排索引，再通过滑动窗口统计词共现，用 PMI 算出关联强度并归一化，自动得到每个词的关联词表。

· 查询扩展：搜索时把原始查询词、手动补充的同义词、自动挖出的关联词按不同权重合并成一个查询向量，让语义相近的词也能命中。

· BM25 打分排序：用 BM25 算法对扩展后的词项逐个计算文档得分，并支持 explain 参数查看每个命中的词和贡献分。

· 内置示例：文件末尾的 main 提供了六个文档和三种查询演示，直接运行即可看到扩展命中效果。

---

优化建议： 您可以调整 build() 中的 window、min_co 和 topk 参数来控制关联词挖掘的范围与数量，也可以把示例文档替换成您自己的真实数据。

文章仅供参考用。