Python+Word2Vec实现中文聊天机器人的示例代码

import bz2
import xml.etree.ElementTree as ET
import re
 
def extract_text(file_path):
    """
    Extract and clean text from a Wikipedia dump file
    """
    with bz2.open(file_path, "r") as f:
        xml = f.read().decode("utf-8")
    root = ET.fromstring("<root>" + xml + "</root>")
    for page in root:
        for revision in page:
            text = revision.find("{http://www.mediawiki.org/xml/export-0.10/}text").text
            clean_text = clean_wiki_text(text)  # Clean text using the clean_wiki_text function
            sentences = split_sentences(clean_text)  # Split cleaned text into sentences using the split_sentences function
            yield from sentences
 
def clean_wiki_text(text):
    """
    Remove markup and other unwanted characters from Wikipedia text
    """
    # Remove markup
    text = re.sub(r"\{\{.*?\}\}", "", text)  # Remove {{...}}
    text = re.sub(r"\[\[.*?\]\]", "", text)  # Remove [...]
    text = re.sub(r"<.*?>", "", text)  # Remove <...>
    text = re.sub(r"&[a-z]+;", "", text)  # Remove &...
    # Remove unwanted characters and leading/trailing white space
    text = text.strip()
    text = re.sub(r"\n+", "\n", text)
    text = re.sub(r"[^\w\s\n!?，。？！]", "", text)  # Remove non-word characters except for !?。.
    text = re.sub(r"\s+", " ", text)
    return text.strip()
 
def split_sentences(text):
    """
    Split text into sentences
    """
    return re.findall(r"[^\n!?。]*[!?。]", text)
 
if __name__ == "__main__":
    file_path = "/path/to/zhwiki-latest-pages-articles.xml.bz2"
    sentences = extract_text(file_path)
    with open("corpus.txt", "w", encoding="utf-8") as f:
        f.write("\n".join(sentences))

import logging
import os.path
import sys
from gensim.corpora import WikiCorpus
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
 
def train_model():
    logging.basicConfig(format="%(asctime)s : %(levelname)s : %(message)s", level=logging.INFO)
 
    input_file = "corpus.txt"
    output_file = "word2vec.model"
 
    # Train Word2Vec model
    sentences = LineSentence(input_file)
    model = Word2Vec(sentences, size=200, window=5, min_count=5, workers=8)
    model.save(output_file)
 
if __name__ == "__main__":
    train_model()

import os
import random
from flask import Flask, request, jsonify
import gensim
 
app = Flask(__name__)
model_file = "word2vec.model"
model = gensim.models.Word2Vec.load(model_file)
chat_log = []
 
@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json()
    input_text = data["input"]
    output_text = get_response(input_text)
    chat_log.append({"input": input_text, "output": output_text})
    return jsonify({"output": output_text})
 
def get_response(input_text):
    # Tokenize input text
    input_tokens = [token for token in jieba.cut(input_text)]
    # Find most similar word in vocabulary
    max_similarity = -1
    best_match = None
    for token in input_tokens:
        if token in model.wv.vocab:
            for match_token in model.wv.most_similar(positive=[token]):
                if match_token[1] > max_similarity:
                    max_similarity = match_token[1]
                    best_match = match_token[0]
    # Generate output text
    if best_match is None:
        return "抱歉，我不知道该如何回答您。"
    else:
        output_text = random.choice([x[0] for x in model.wv.most_similar(positive=[best_match])])
        return output_text
 
if __name__ == "__main__":
    app.run(debug=True)