Upload 2 files

macro-correct.py

@@ -0,0 +1,114 @@
+# !/usr/bin/python
+# -*- coding: utf-8 -*-
+# @time    : 2021/2/29 21:41
+# @author  : Mo
+# @function: transformers直接加载bert类模型测试
+
+
+import traceback
+import time
+import sys
+import os
+os.environ["USE_TORCH"] = "1"
+
+from transformers import BertConfig, BertTokenizer, BertForMaskedLM
+import gradio as gr
+import torch
+
+
+
+# pretrained_model_name_or_path = "shibing624/macbert4csc-base-chinese"
+pretrained_model_name_or_path = "Macropodus/macbert4mdcspell_v1"
+# pretrained_model_name_or_path = "Macropodus/macbert4csc_v1"
+# pretrained_model_name_or_path = "Macropodus/macbert4csc_v2"
+# pretrained_model_name_or_path = "Macropodus/bert4csc_v1"
+device = torch.device("cpu")
+# device = torch.device("cuda")
+max_len = 128
+
+print("load model, please wait a few minute!")
+tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path)
+bert_config = BertConfig.from_pretrained(pretrained_model_name_or_path)
+model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path)
+model.to(device)
+print("load model success!")
+
+texts = [
+    "机七学习是人工智能领遇最能体现智能的一个分知",
+    "我是练习时长两念半的鸽仁练习生蔡徐坤",
+]
+len_mid = min(max_len, max([len(t)+2 for t in texts]))
+
+with torch.no_grad():
+    outputs = model(**tokenizer(texts, padding=True, max_length=len_mid,
+                                return_tensors="pt").to(device))
+
+def get_errors(source, target):
+    """   极简方法获取 errors   """
+    len_min = min(len(source), len(target))
+    errors = []
+    for idx in range(len_min):
+        if source[idx] != target[idx]:
+            errors.append([source[idx], target[idx], idx])
+    return errors
+
+result = []
+for probs, source in zip(outputs.logits, texts):
+    ids = torch.argmax(probs, dim=-1)
+    tokens_space = tokenizer.decode(ids[1:-1], skip_special_tokens=False)
+    text_new = tokens_space.replace(" ", "")
+    target = text_new[:len(source)]
+    errors = get_errors(source, target)
+    print(source, " => ", target, errors)
+    result.append([target, errors])
+print(result)
+
+
+def macro_correct(text):
+    with torch.no_grad():
+        outputs = model(**tokenizer([text], padding=True, max_length=max_len,
+                                    return_tensors="pt").to(device))
+
+    def to_highlight(corrected_sent, errs):
+        output = [{"entity": "纠错", "word": err[1], "start": err[2], "end": err[3]} for i, err in
+                  enumerate(errs)]
+        return {"text": corrected_sent, "entities": output}
+
+    def get_errors(source, target):
+        """   极简方法获取 errors   """
+        len_min = min(len(source), len(target))
+        errors = []
+        for idx in range(len_min):
+            if source[idx] != target[idx]:
+                errors.append([source[idx], target[idx], idx])
+        return errors
+
+    result = []
+    for probs, source in zip(outputs.logits, texts):
+        ids = torch.argmax(probs, dim=-1)
+        tokens_space = tokenizer.decode(ids[1:-1], skip_special_tokens=False)
+        text_new = tokens_space.replace(" ", "")
+        target = text_new[:len(source)]
+        errors = get_errors(source, target)
+        print(source, " => ", target, errors)
+        result.append([target, errors])
+        # print(result)
+        return target + " " + str(errors)
+
+
+if __name__ == '__main__':
+    print(macro_correct('少先队员因该为老人让坐'))
+
+    text_sample = '机七学习是人工智能领遇最能体现智能的一个分知'
+
+    gr.Interface(
+        macro_correct,
+        inputs='text',
+        outputs='text',
+        title="Chinese Spelling Correction Model Macropodus/macbert4csc_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=text_sample
+    ).launch()
+    # ).launch(server_name="0.0.0.0", server_port=8066, share=False, debug=True)
+

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+transformers
+torch