Python基于BERT模型实现上下文纠错功能详解

pip install transformers torch

from transformers import BertTokenizer, BertForMaskedLM
import torch

# 初始化tokenizer和model
model_name = "bert-base-chinese"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForMaskedLM.from_pretrained(model_name)
model.eval()  # 切换至评估模式

def predict_masked_word(text, top_k=5):
    """
    输入带[MASK]的文本，返回top-k最可能的词及其概率
    """
    # 编码输入
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]

    # 模型推理
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        mask_logits = logits[0, mask_token_index, :]
        probs = torch.softmax(mask_logits, dim=-1)

    # 取top-k结果
    top_results = torch.topk(probs, top_k)
    predicted_tokens = [tokenizer.decode([token_id]) for token_id in top_results.indices[0]]
    scores = top_results.values[0].tolist()

    return list(zip(predicted_tokens, scores))

text = "春风又绿江南[MASK]"
results = predict_masked_word(text)
for word, score in results:
    print(f"{word}({score:.1%})")

def multi_mask_correct(text):
    # 示例：处理“他再家[MASK]作业”
    step1_text = "他[MASK]家写作业"  # 先修正“再”→“在”
    step1_results = predict_masked_word(step1_text)
    print("第一步修正:", step1_results)

    step2_text = "他在家[MASK]作业"  # 再修正“[MASK]”→“写”
    step2_results = predict_masked_word(step2_text)
    print("第二步修正:", step2_results)

    return "他在家写作业"

def apply_rules(text, predictions):
    filtered_results = []
    for word, score in predictions:
        if len(word) == 1 or word.isalpha():  # 示例规则：单字或纯字母
            filtered_results.append((word, score))
    return filtered_results

text = "我喜换吃苹果"
masked_text = "我喜[MASK]吃苹果"
predictions = predict_masked_word(masked_text)
filtered_results = apply_rules("喜", predictions)  # 假设“喜”为错误位置
print("过滤后结果:", filtered_results)

from transformers import BertForMaskedLM, BertConfig

# 加载基础模型
config = BertConfig.from_pretrained(model_name)
model = BertForMaskedLM(config)

# 继续预训练（需准备领域语料）
# model.train()
# for epoch in range(3):
#     train_loop(model, domain_data)

def writing_assistant(text):
    # 示例：检测“的得地”混淆
    if "的" in text and "地" in text:
        # 调用BERT模型进一步分析
        pass
    return corrected_text

def optimize_reply(reply):
    masked_reply = reply.replace("嘛", "[MASK]")  # 示例：修正“嘛”→“吗”
    predictions = predict_masked_word(masked_reply)
    if predictions[0][0] == "吗":
        return reply.replace("嘛", "吗")
    return reply

def ocr_correction(ocr_text, image_context):
    # 示例：根据图像中的“银行”标志修正OCR错误“邮局”
    if "邮局" in ocr_text and "银行" in image_context:
        masked_text = ocr_text.replace("邮局", "[MASK]")
        predictions = predict_masked_word(masked_text)
        if "银行" in [p[0] for p in predictions]:
            return ocr_text.replace("邮局", "银行")
    return ocr_text