[update]edit main

.gitignore

@@ -5,5 +5,6 @@
 **/flagged/
 **/__pycache__/
 
+cache/
 flagged/
 trained_models/

main.py

@@ -4,6 +4,11 @@ import argparse
 from collections import defaultdict
 import os
 import platform
+import re
+
+from project_settings import project_path
+
+os.environ["HUGGINGFACE_HUB_CACHE"] = (project_path / "cache/huggingface/hub").as_posix()
 
 import gradio as gr
 from threading import Thread
@@ -12,8 +17,6 @@ from transformers.models.bert.tokenization_bert import BertTokenizer
 from transformers.generation.streamers import TextIteratorStreamer
 import torch
 
-from project_settings import project_path
-
 
 def get_args():
     parser = argparse.ArgumentParser()
@@ -38,6 +41,11 @@ examples = [
 ]
 
 
+def repl(match):
+    result = "{}{}".format(match.group(1), match.group(2))
+    return result
+
+
 def main():
     args = get_args()
 
@@ -94,13 +102,16 @@ def main():
             if first_answer:
                 first_answer = False
                 continue
-            # output_ = output_.replace(text, "")
-            # output_ = output_.replace("[CLS]", "")
-            output_ = output_.replace("[SEP]", "\n")
+
+            output_ = output_.replace("[UNK] ", "")
             output_ = output_.replace("[UNK]", "")
-            output_ = output_.replace(" ", "")
 
-            output += output_.strip()
+            output += output_
+
+            output = output.lstrip("[SEP] ,.!?")
+            output = output.replace("[SEP]", "\n")
+            output = re.sub(r"([\u4e00-\u9fa5]) ([\u4e00-\u9fa5])", repl, output)
+
             output_text_box.value += output
             yield output