Add small100 to pythainlp.translate.Translate

Author: wannaphong

pythainlp/translate/core.py

@@ -41,19 +41,26 @@ class Translate:
     """
 
     def __init__(
-        self, src_lang: str, target_lang: str, use_gpu: bool = False
+        self, src_lang: str, target_lang: str, engine: str="default", use_gpu: bool = False
     ) -> None:
         """
         :param str src_lang: source language
         :param str target_lang: target language
+        :param str engine: Machine Translation engine
         :param bool use_gpu: load model to gpu (Default is False)
 
+        **Options for engine*
+            * *default* - The engine default by each a language.
+            * *small100* - A multilingual machine translation model (covering 100 languages)
+
         **Options for source & target language**
             * *th* - *en* - Thai to English
             * *en* - *th* - English to Thai
             * *th* - *zh* - Thai to Chinese
             * *zh* - *th* - Chinese to Thai
             * *th* - *fr* - Thai to French
+            * *th* - *xx* - Thai to xx (xx is language code). It uses small100 model.
+            * *xx* - *th* - xx to Thai (xx is language code). It uses small100 model.
 
         :Example:
 
@@ -66,10 +73,21 @@ def __init__(
             # output: I love cat.
         """
         self.model = None
-        self.load_model(src_lang, target_lang, use_gpu)
-
-    def load_model(self, src_lang: str, target_lang: str, use_gpu: bool):
-        if src_lang == "th" and target_lang == "en":
+        self.engine = engine
+        self.src_lang = src_lang
+        self.use_gpu = use_gpu
+        self.target_lang = target_lang
+        self.load_model()
+
+    def load_model(self):
+        src_lang = self.src_lang
+        target_lang = self.target_lang
+        use_gpu = self.use_gpu
+        if self.engine == "small100":
+            from .small100 import Small100Translator
+
+            self.model = Small100Translator(use_gpu)
+        elif src_lang == "th" and target_lang == "en":
             from pythainlp.translate.en_th import ThEnTranslator
 
             self.model = ThEnTranslator(use_gpu)
@@ -100,4 +118,6 @@ def translate(self, text) -> str:
         :return: translated text in target language
         :rtype: str
         """
+        if self.engine == "small100":
+            return self.model.translate(text, tgt_lang=self.target_lang)
         return self.model.translate(text)

tests/test_translate.py

@@ -73,5 +73,17 @@ def test_translate(self):
                 "ทดสอบระบบ",
             )
         )
+        self.th_fr_translator = Translate('th', 'fr', engine="small100")
+        self.assertIsNotNone(
+            self.th_fr_translator.translate(
+                "ทดสอบระบบ",
+            )
+        )
+        self.th_fr_translator = Translate('th', 'ja')
+        self.assertIsNotNone(
+            self.th_fr_translator.translate(
+                "ทดสอบระบบ",
+            )
+        )
         with self.assertRaises(ValueError):
-            self.th_cat_translator = Translate('th', 'cat')
+            self.th_cat_translator = Translate('th', 'cat', engine="fkfj")