add more opt position and bpe helper (facebookresearch#2550)

tancnle · Apr 13, 2020 · d694b99 · d694b99
1 parent 0d9e16d
commit d694b99
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/parlai/utils/pickle.py b/parlai/utils/pickle.py
@@ -33,15 +33,22 @@ class Unpickler(pickle._Unpickler):  # type: ignore
     def find_class(self, module, name):
         try:
             return super().find_class(module, name)
-        except ModuleNotFoundError:
+        except (ModuleNotFoundError, AttributeError):
             if module.startswith('apex.'):
                 # user doesn't have apex installed. We'll deal with this later.
                 return FakeAPEXClass
             else:
-                if module == 'parlai.core.utils' and name == 'Opt':
+                if (
+                    module == 'parlai.core.utils' or module == 'parlai.utils.misc'
+                ) and name == 'Opt':
                     from parlai.core.opt import Opt
 
                     return Opt
+                if module == 'parlai.core.dict' and name == '_BPEHelper':
+                    from parlai.utils.bpe import SubwordBPEHelper as _BPEHelper
+
+                    return _BPEHelper
+
                 raise