Create bart_pipeline.py

bart_pipeline.py

@@ -0,0 +1,99 @@
+from transformers import TranslationPipeline
+from transformers.pipelines.text2text_generation import ReturnType
+from transformers import BartForConditionalGeneration, BertTokenizer
+import logging
+import re
+
+
+def fix_chinese_text_generation_space(text):
+    output_text = text
+    output_text = re.sub(
+        r'([\u3401-\u9FFF+——！，。？、~@#￥%…&*（）：；《）《》“”()»〔〕\-!$^*()_+|~=`{}\[\]:";\'<>?,.\/\\])\s([^0-9a-zA-Z])', r'\1\2', output_text)
+    output_text = re.sub(
+        r'([^0-9a-zA-Z])\s([\u3401-\u9FFF+——！，。？、~@#￥%…&*（）：；《）《》“”()»〔〕\-!$^*()_+|~=`{}\[\]:";\'<>?,.\/\\])', r'\1\2', output_text)
+    output_text = re.sub(
+        r'([\u3401-\u9FFF+——！，。？、~@#￥%…&*（）：；《）《》“”()»〔〕\-!$^*()_+|~=`{}\[\]:";\'<>?,.\/\\])\s([a-zA-Z0-9])', r'\1\2', output_text)
+    output_text = re.sub(
+        r'([a-zA-Z0-9])\s([\u3401-\u9FFF+——！，。？、~@#￥%…&*（）：；《）《》“”()»〔〕\-!$^*()_+|~=`{}\[\]:";\'<>?,.\/\\])', r'\1\2', output_text)
+    output_text = re.sub(r'＄\s([0-9])', r'＄\1', output_text)
+    output_text = re.sub(',', '，', output_text)
+    output_text = re.sub(r'([0-9])，([0-9])', r'\1,\2',
+                         output_text)  # fix comma in numbers
+    # fix multiple commas
+    output_text = re.sub(r'\s?[，]+\s?', '，', output_text)
+    output_text = re.sub(r'\s?[、]+\s?', '、', output_text)
+    # fix period
+    output_text = re.sub(r'\s?[。]+\s?', '。', output_text)
+    # fix ...
+    output_text = re.sub(r'\s?\.{3,}\s?', '...', output_text)
+    # fix exclamation mark
+    output_text = re.sub(r'\s?[!！]+\s?', '！', output_text)
+    # fix question mark
+    output_text = re.sub(r'\s?[?？]+\s?', '？', output_text)
+    # fix colon
+    output_text = re.sub(r'\s?[:：]+\s?', '：', output_text)
+    # fix quotation mark
+    output_text = re.sub(r'\s?(["“”\']+)\s?', r'\1', output_text)
+    # fix semicolon
+    output_text = re.sub(r'\s?[;；]+\s?', '；', output_text)
+    # fix dots
+    output_text = re.sub(r'\s?([~●．…]+)\s?', r'\1', output_text)
+    output_text = re.sub(r'\s?\[…\]\s?', '', output_text)
+    output_text = re.sub(r'\s?\[\.\.\.\]\s?', '', output_text)
+    output_text = re.sub(r'\s?\.{3,}\s?', '...', output_text)
+    # fix slash
+    output_text = re.sub(r'\s?[／/]+\s?', '／', output_text)
+    # fix dollar sign
+    output_text = re.sub(r'\s?[$＄]+\s?', '＄', output_text)
+    # fix @
+    output_text = re.sub(r'\s?([@＠]+)\s?', '＠', output_text)
+    # fix baskets
+    output_text = re.sub(
+        r'\s?([\[\(<〖【「『（）』」】〗>\)\]]+)\s?', r'\1', output_text)
+
+    return output_text
+
+
+class BartPipeline(TranslationPipeline):
+    def __init__(self,
+                 model_name_or_path: str = "indiejoseph/bart-base-cantonese",
+                 device=None,
+                 max_length=512,
+                 src_lang=None,
+                 tgt_lang=None):
+        self.model_name_or_path = model_name_or_path
+        self.tokenizer = self._load_tokenizer()
+        self.model = self._load_model()
+        self.model.eval()
+        super().__init__(self.model, self.tokenizer, device=device,
+                         max_length=max_length, src_lang=src_lang, tgt_lang=tgt_lang)
+
+    def _load_tokenizer(self):
+        return BertTokenizer.from_pretrained(self.model_name_or_path)
+
+    def _load_model(self):
+        return BartForConditionalGeneration.from_pretrained(self.model_name_or_path)
+
+    def postprocess(
+        self,
+        model_outputs,
+        return_type=ReturnType.TEXT,
+        clean_up_tokenization_spaces=True,
+    ):
+        records = super().postprocess(
+            model_outputs,
+            return_type=return_type,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+        )
+        for rec in records:
+            translation_text = fix_chinese_text_generation_space(
+                rec["translation_text"].strip())
+
+            rec["translation_text"] = translation_text
+        return records
+
+
+if __name__ == '__main__':
+    pipe = BartPipeline(device=0)
+
+    print(pipe('哈哈，我正在努力研究緊個問題。不過，邊個知呢，可能哪一日我會諗到一個好主意去實現到佢。', max_length=100))