Update app.py

app.py

@@ -12,6 +12,13 @@ model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path)
 vocab = tokenizer.vocab
 
 
+# from modelscope import AutoTokenizer, AutoModelForMaskedLM
+# pretrained_model_name_or_path = "Macadam/macbert4mdcspell_v2"
+# tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
+# model = AutoModelForMaskedLM.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'))
@@ -29,7 +36,53 @@ def func_macro_correct(text):
                 return False
         return True
 
-    def get_errors(corrected_text, origin_text, unk_tokens=[], know_tokens=[]):
+    def get_errors_from_diff_length(corrected_text, origin_text, unk_tokens=[], know_tokens=[]):
+        """Get errors between corrected text and origin text
+        code from:  https://github.com/shibing624/pycorrector
+        """
+        new_corrected_text = ""
+        errors = []
+        i, j = 0, 0
+        unk_tokens = unk_tokens or [' ', '“', '”', '‘', '’', '琊', '\n', '…', '擤', '\t', '玕', '']
+        while i < len(origin_text) and j < len(corrected_text):
+            if origin_text[i] in unk_tokens or origin_text[i] not in know_tokens:
+                new_corrected_text += origin_text[i]
+                i += 1
+            elif corrected_text[j] in unk_tokens:
+                new_corrected_text += corrected_text[j]
+                j += 1
+            # Deal with Chinese characters
+            elif flag_total_chinese(origin_text[i]) and flag_total_chinese(corrected_text[j]):
+                # If the two characters are the same, then the two pointers move forward together
+                if origin_text[i] == corrected_text[j]:
+                    new_corrected_text += corrected_text[j]
+                    i += 1
+                    j += 1
+                else:
+                    # Check for insertion errors
+                    if j + 1 < len(corrected_text) and origin_text[i] == corrected_text[j + 1]:
+                        errors.append(('', corrected_text[j], j))
+                        new_corrected_text += corrected_text[j]
+                        j += 1
+                    # Check for deletion errors
+                    elif i + 1 < len(origin_text) and origin_text[i + 1] == corrected_text[j]:
+                        errors.append((origin_text[i], '', i))
+                        i += 1
+                    # Check for replacement errors
+                    else:
+                        errors.append((origin_text[i], corrected_text[j], i))
+                        new_corrected_text += corrected_text[j]
+                        i += 1
+                        j += 1
+            else:
+                new_corrected_text += origin_text[i]
+                if origin_text[i] == corrected_text[j]:
+                    j += 1
+                i += 1
+        errors = sorted(errors, key=operator.itemgetter(2))
+        return new_corrected_text, errors
+
+    def get_errors_from_same_length(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
         """
@@ -57,7 +110,14 @@ def func_macro_correct(text):
 
     _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("#" * 128)
+    print(text)
+    print(corrected_text)
+    print(len(text), len(corrected_text))
+    if len(corrected_text) == len(text):
+        corrected_text, details = get_errors_from_same_length(corrected_text, text, know_tokens=vocab)
+    else:
+        corrected_text, details = get_errors_from_diff_length(corrected_text, text, know_tokens=vocab)
     print(text, ' => ', corrected_text, details)
     return corrected_text + ' ' + str(details)
 
@@ -88,6 +148,6 @@ if __name__ == '__main__':
         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()
+    ).launch()  # .launch(server_name="0.0.0.0", server_port=8036, share=False, debug=True)