Python实战教程：基于维基百科数据的开放域问答机器人DrQA实现与效果展示

我最近在搞一个基于维基百科的问答机器人，用的就是Facebook Research开源的DrQA框架。这玩意儿挺有意思的，核心就是把信息检索和阅读理解结合起来。

简单说下实现思路：先用文档检索器从维基百科里找到相关文章，然后用阅读理解模型从文章里提取答案。我用的TF-IDF做检索，阅读理解模型是预训练的BERT。

import wikipedia
from rank_bm25 import BM25Okapi
from transformers import BertTokenizer, BertForQuestionAnswering
import torch

class SimpleDrQA:
    def __init__(self):
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        self.model = BertForQuestionAnswering.from_pretrained('bert-base-uncased')
        self.documents = []
        
    def fetch_wiki_data(self, query, num_pages=3):
        """从维基百科获取相关页面"""
        search_results = wikipedia.search(query, results=num_pages)
        for title in search_results:
            try:
                page = wikipedia.page(title, auto_suggest=False)
                self.documents.append(page.content)
            except:
                continue
        
    def retrieve_documents(self, question, top_k=3):
        """用BM25检索相关文档段落"""
        tokenized_docs = [doc.split() for doc in self.documents]
        bm25 = BM25Okapi(tokenized_docs)
        tokenized_query = question.split()
        scores = bm25.get_scores(tokenized_query)
        top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
        return [self.documents[i] for i in top_indices]
    
    def extract_answer(self, question, context):
        """用BERT模型提取答案"""
        inputs = self.tokenizer.encode_plus(
            question, context, 
            add_special_tokens=True, 
            return_tensors='pt',
            max_length=512,
            truncation=True
        )
        
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        start_scores = outputs.start_logits
        end_scores = outputs.end_logits
        
        start_idx = torch.argmax(start_scores)
        end_idx = torch.argmax(end_scores) + 1
        
        answer_tokens = inputs['input_ids'][0][start_idx:end_idx]
        answer = self.tokenizer.decode(answer_tokens)
        
        return answer
    
    def ask(self, question):
        """主问答接口"""
        # 获取维基数据
        self.fetch_wiki_data(question)
        
        # 检索相关文档
        contexts = self.retrieve_documents(question)
        
        # 从最佳文档中提取答案
        best_answer = ""
        best_score = -1
        
        for context in contexts:
            answer = self.extract_answer(question, context[:2000])  # 限制上下文长度
            # 这里可以加个置信度评分
            if answer and len(answer) > 5:  # 简单过滤短答案
                best_answer = answer
                break
                
        return best_answer

# 使用示例
if __name__ == "__main__":
    qa = SimpleDrQA()
    question = "Who invented Python programming language?"
    answer = qa.ask(question)
    print(f"Q: {question}")
    print(f"A: {answer}")

跑起来效果还行，对事实类问题回答得比较准。比如问"Python谁发明的"，它能从维基百科里找到Guido van Rossum。不过处理复杂问题或者需要推理的时候就不太行了。

这只是一个简化版，真正的DrQA系统要复杂得多，包括更好的文档处理、多段落答案融合、置信度校准这些。但基本框架就是这样：检索+阅读理解。

建议试试用更先进的检索模型和更大的预训练模型来提升效果。