app.py

# -*- coding: utf-8 -*-

import gradio as gr
import operator
import torch
from transformers import BertTokenizer, BertForMaskedLM


pretrained_model_name_or_path = "Macropodus/macbert4mdcspell_v2"
tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path)
model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path)
vocab = tokenizer.vocab


def func_macro_correct(text):
    with torch.no_grad():
        outputs = model(**tokenizer([text], padding=True, return_tensors='pt'))

    def flag_total_chinese(text):
        """
        judge is total chinese or not, 判断是不是全是中文
        Args:
            text: str, eg. "macadam, 碎石路"
        Returns:
            bool, True or False
        """
        for word in text:
            if not "\u4e00" <= word <= "\u9fa5":
                return False
        return True

    def get_errors(corrected_text, origin_text, unk_tokens=[], know_tokens=[]):
        """Get new corrected text and errors between corrected text and origin text
        code from:  https://github.com/shibing624/pycorrector
        """
        errors = []
        unk_tokens = unk_tokens or [' ', '“', '”', '‘', '’', '琊', '\n', '…', '擤', '\t', '玕', '', '，']

        for i, ori_char in enumerate(origin_text):
            if i >= len(corrected_text):
                continue
            if ori_char in unk_tokens or ori_char not in know_tokens:
                # deal with unk word
                corrected_text = corrected_text[:i] + ori_char + corrected_text[i + 1:]
                continue
            if ori_char != corrected_text[i]:
                if not flag_total_chinese(ori_char):
                    # pass not chinese char
                    corrected_text = corrected_text[:i] + ori_char + corrected_text[i + 1:]
                    continue
                if not flag_total_chinese(corrected_text[i]):
                    corrected_text = corrected_text[:i] + corrected_text[i + 1:]
                    continue
                errors.append([ori_char, corrected_text[i], i])
        errors = sorted(errors, key=operator.itemgetter(2))
        return corrected_text, errors

    _text = tokenizer.decode(torch.argmax(outputs.logits[0], dim=-1), skip_special_tokens=True).replace(' ', '')
    corrected_text = _text[:len(text)]
    corrected_text, details = get_errors(corrected_text, text, know_tokens=vocab)
    print(text, ' => ', corrected_text, details)
    return corrected_text + ' ' + str(details)


if __name__ == '__main__':
    print(func_macro_correct('他法语说的很好，的语也不错'))

    examples = [
        "夫谷之雨，犹复云之亦从的起，因与疾风俱飘，参于天，集于的。",
        "机七学习是人工智能领遇最能体现智能的一个分知",
        '他们的吵翻很不错，再说他们做的咖喱鸡也好吃',
        "抗疫路上，除了提心吊胆也有难的得欢笑。",
        "我是练习时长两念半的鸽仁练习生蔡徐坤",
        "清晨，如纱一般地薄雾笼罩着世界。",
        "得府许我立庙于此，故请君移去尔。",
        "他法语说的很好，的语也不错",
        "遇到一位很棒的奴生跟我疗天",
        "五年级得数学，我考的很差。",
        "我们为这个目标努力不解",
        '今天兴情很好',
    ]

    gr.Interface(
        func_macro_correct,
        inputs='text',
        outputs='text',
        title="Chinese Spelling Correction Model Macropodus/macbert4mdcspell_v2",
        description="Copy or input error Chinese text. Submit and the machine will correct text.",
        article="Link to <a href='https://github.com/yongzhuo/macro-correct' style='color:blue;' target='_blank\'>Github REPO: macro-correct</a>",
        examples=examples
    ).launch(server_name="0.0.0.0", server_port=8036, share=False, debug=True)