Add small100

Author: wannaphong

pythainlp/translate/small100.py

@@ -0,0 +1,60 @@
+from transformers import M2M100ForConditionalGeneration
+from .tokenization_small100 import SMALL100Tokenizer
+
+class Small100Translator:
+    """
+    Machine Translation with small100 model
+
+    - Huggingface https://huggingface.co/alirezamsh/small100
+
+    :param bool use_gpu : load model to gpu (Default is False)
+    """
+
+    def __init__(
+        self,
+        use_gpu: bool = False,
+        pretrained: str = "alirezamsh/small100",
+    ) -> None:
+        self.pretrained = pretrained
+        self.model = M2M100ForConditionalGeneration.from_pretrained(self.pretrained)
+        self.tgt_lang = None
+        if use_gpu:
+            self.model = self.model.cuda()
+
+    def translate(self, text: str, tgt_lang: str="en") -> str:
+        """
+        Translate text from X to X
+
+        :param str text: input text in source language
+        :param str tgt_lang: target language
+        :return: translated text in target language
+        :rtype: str
+
+        :Example:
+
+        ::
+
+            from pythainlp.translate.small100 import Small100Translator
+
+            mt = Small100Translator()
+
+            # Translate text from Thai to English
+            mt.translate("ทดสอบระบบ", tgt_lang="en")
+            # output: 'Testing system'
+
+            # Translate text from Thai to Chinese
+            mt.translate("ทดสอบระบบ", tgt_lang="zh")
+            # output: '系统测试'
+
+            # Translate text from Thai to French
+            mt.translate("ทดสอบระบบ", tgt_lang="fr")
+            # output: 'Test du système'
+
+        """
+        if tgt_lang!=self.tgt_lang:
+            self.tokenizer = SMALL100Tokenizer.from_pretrained(self.pretrained, tgt_lang=tgt_lang)
+            self.tgt_lang = tgt_lang
+        self.translated = self.model.generate(
+            **self.tokenizer(text, return_tensors="pt")
+        )
+        return self.tokenizer.batch_decode(self.translated, skip_special_tokens=True)[0]