feat(base): support loading external modules from py and yaml

gnes-ai · Jul 19, 2019 · 66f26e0 · 66f26e0
1 parent 07ce8e9
commit 66f26e0
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 4 deletions.
diff --git a/tests/contrib/dummy.yml b/tests/contrib/dummy.yml
@@ -0,0 +1,3 @@
+!FooContribEncoder
+parameter:
+  bar: 531
diff --git a/tests/contrib/dummy_contrib.py b/tests/contrib/dummy_contrib.py
@@ -0,0 +1,12 @@
+from gnes.encoder.base import BaseTextEncoder
+
+
+class FooContribEncoder(BaseTextEncoder):
+
+    def __init__(self, bar: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_trained = True
+        self.bar = bar
+
+    def encode(self, text, **kwargs):
+        return 'hello %d' % self.bar
diff --git a/tests/test_contrib_module.py b/tests/test_contrib_module.py
@@ -6,11 +6,11 @@ class TestYaml(unittest.TestCase):
 
     def setUp(self):
         dirname = os.path.dirname(__file__)
-        module_path = os.path.join(dirname, 'contrib', 'toy22.py')
+        module_path = os.path.join(dirname, 'contrib', 'dummy_contrib.py')
         cls_name = 'FooContribEncoder'
         os.environ['GNES_CONTRIB_MODULE'] = '%s:%s' % (cls_name, module_path)
         self.yaml_path = os.path.join(os.path.dirname(__file__),
-                                      'contrib', 'toy22.yml')
+                                      'contrib', 'dummy.yml')
 
     def test_load_contrib(self):
         from gnes.encoder.base import BaseEncoder, BaseTextEncoder

diff --git a/tests/test_contrib_module_negative.py b/tests/test_contrib_module_negative.py
@@ -8,11 +8,11 @@ class TestYaml(unittest.TestCase):
 
     def setUp(self):
         dirname = os.path.dirname(__file__)
-        module_path = os.path.join(dirname, 'contrib', 'toy22.py')
+        module_path = os.path.join(dirname, 'contrib', 'dummy_contrib.py')
         cls_name = 'FooContribEncoder'
         os.environ['GNES_CONTRIB_MODULE'] = '%s:%s' % (cls_name, module_path)
         self.yaml_path = os.path.join(os.path.dirname(__file__),
-                                      'contrib', 'toy22.yml')
+                                      'contrib', 'dummy.yml')
 
     def test_broken_contrib(self):
         os.environ['GNES_CONTRIB_MODULE'] = ''