基于Unigram和Bigram构建语料库倒排索引并实现查询处理

需求描述 #

编写代码遍历语料库，创建包含每个文档中Unigram（一元语法）和Bigram（二元语法）的倒排索引，并实现对用户输入查询的响应功能。

现有实现代码 #

import os
import re
from collections import defaultdict
import time

corpus_directory = "corpus path"

def create_inverted_index(corpus_directory, encoding='utf-8'):
    inverted_index = {}
    for root, dirs, files in os.walk(corpus_directory):
        for file in files:
            file_path = os.path.join(root, file)
            with open(file_path, 'r', encoding=encoding, errors='ignore') as f:
                document = f.read()
                tokens = tokenize_text(document, lowercase=True)
                bigrams = [" ".join(tokens[i:i+2]) for i in range(len(tokens) - 1)]
                all_tokens = tokens + bigrams
                for token in all_tokens:
                    if token not in inverted_index:
                        inverted_index[token] = []
                    inverted_index[token].append(file_path)
    return inverted_index

当前状态 #

已完成包含各文档Unigram和Bigram的倒排索引创建，现需实现对用户输入查询的响应功能。

补充实现内容 #

1. 完善分词函数 #

原代码调用了tokenize_text但未定义，先补充该基础分词逻辑：

def tokenize_text(text, lowercase=True):
    if lowercase:
        text = text.lower()
    # 提取字母数字作为分词单元，可根据需求调整正则规则适配中文或特殊术语
    tokens = re.findall(r'\w+', text)
    return tokens

2. 实现查询响应函数 #

将用户查询分词生成Unigram和Bigram，从倒排索引中匹配相关文档并返回去重结果：

def respond_to_query(query, inverted_index, lowercase=True):
    query_tokens = tokenize_text(query, lowercase)
    query_bigrams = [" ".join(query_tokens[i:i+2]) for i in range(len(query_tokens) - 1)]
    query_terms = query_tokens + query_bigrams
    
    matched_docs = set()
    for term in query_terms:
        if term in inverted_index:
            matched_docs.update(inverted_index[term])
    
    return sorted(matched_docs)

3. 完整使用示例 #

if __name__ == "__main__":
    # 创建倒排索引
    start_time = time.time()
    index = create_inverted_index(corpus_directory)
    print(f"倒排索引创建完成，耗时 {time.time() - start_time:.2f} 秒")
    
    # 处理用户查询
    user_query = input("请输入查询内容：")
    results = respond_to_query(user_query, index)
    
    if results:
        print(f"找到 {len(results)} 篇相关文档：")
        for doc in results:
            print(f"- {doc}")
    else:
        print("未找到相关文档")

说明 #

• 分词规则可根据语料类型调整，比如适配中文分词可替换为jieba等专业库
• 默认采用并集匹配逻辑：文档包含任意查询词（Unigram/Bigram）即被返回；若需精准匹配，可修改为取所有查询词对应文档集合的交集
• 用集合存储匹配文档实现自动去重，避免重复输出

内容的提问来源于stack exchange，提问作者Shourya Manekar