From 4e4297df73e88b2ff576a59d7b7ed38121903d63 Mon Sep 17 00:00:00 2001
From: wuzewu <wuzewu@baidu.com>
Date: Mon, 1 Apr 2019 11:18:39 +0800
Subject: [PATCH] support init module with key

---
 paddle_hub/__init__.py         |  1 +
 paddle_hub/dataset/__init__.py |  2 ++
 paddle_hub/module/manager.py   |  9 ++++++--
 paddle_hub/module/module.py    | 42 +++++++++++++++++++++++++---------
 4 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/paddle_hub/__init__.py b/paddle_hub/__init__.py
index 1fddbb74..2031412b 100644
--- a/paddle_hub/__init__.py
+++ b/paddle_hub/__init__.py
@@ -37,3 +37,4 @@ from .finetune.config import FinetuneConfig
 from .finetune.task import Task
 
 from .reader import BERTTokenizeReader
+from .reader.cv_reader import ImageClassificationReader
diff --git a/paddle_hub/dataset/__init__.py b/paddle_hub/dataset/__init__.py
index 5d2d4674..b059666e 100644
--- a/paddle_hub/dataset/__init__.py
+++ b/paddle_hub/dataset/__init__.py
@@ -15,3 +15,5 @@
 from .dataset import InputExample, HubDataset
 from .chnsenticorp import ChnSentiCorp
 from .msra_ner import MSRA_NER
+from .dogcat import DogCatDataset as DogCat
+from .flowers import FlowersDataset as Flowers
diff --git a/paddle_hub/module/manager.py b/paddle_hub/module/manager.py
index 8e530da2..06069a5c 100644
--- a/paddle_hub/module/manager.py
+++ b/paddle_hub/module/manager.py
@@ -21,12 +21,13 @@ import shutil
 
 from paddle_hub.common import utils
 from paddle_hub.common.downloader import default_downloader
+from paddle_hub.common.dir import MODULE_HOME
 import paddle_hub as hub
 
 
 class LocalModuleManager:
     def __init__(self, module_home=None):
-        self.local_modules_dir = module_home if module_home else hub.MODULE_HOME
+        self.local_modules_dir = module_home if module_home else MODULE_HOME
         self.modules_dict = {}
         if not os.path.exists(self.local_modules_dir):
             utils.mkdir(self.local_modules_dir)
@@ -77,7 +78,11 @@ class LocalModuleManager:
             save_name=module_name,
             replace=True)
         result, tips, module_dir = default_downloader.uncompress(
-            file=module_zip_file, dirname=hub.MODULE_HOME, delete_file=True)
+            file=module_zip_file, dirname=MODULE_HOME, delete_file=True)
+
+        save_path = os.path.join(MODULE_HOME, module_name)
+        shutil.move(module_dir, save_path)
+        module_dir = save_path
 
         if module_dir:
             tips = "Successfully installed %s" % module_name
diff --git a/paddle_hub/module/module.py b/paddle_hub/module/module.py
index 5616afc9..7ba43394 100644
--- a/paddle_hub/module/module.py
+++ b/paddle_hub/module/module.py
@@ -15,24 +15,28 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+
+import os
+import time
+import sys
+import functools
+from shutil import copyfile
+
+import paddle
+import paddle.fluid as fluid
+
 from paddle_hub.common import utils
+from paddle_hub.common import paddle_helper
 from paddle_hub.common.logger import logger
 from paddle_hub.common.downloader import default_downloader
-from paddle_hub.common import paddle_helper
 from paddle_hub.module import module_desc_pb2
 from paddle_hub.module import check_info_pb2
 from paddle_hub.module.signature import Signature, create_signature
 from paddle_hub.module.checker import ModuleChecker
+from paddle_hub.module.manager import default_module_manager
+from paddle_hub.module.base_processor import BaseProcessor
 from paddle_hub.io.reader import yaml_reader
 from paddle_hub import version
-from paddle_hub.module.base_processor import BaseProcessor
-from shutil import copyfile
-import os
-import time
-import sys
-import functools
-import paddle
-import paddle.fluid as fluid
 
 __all__ = ['Module', 'create_module']
 
@@ -89,6 +93,7 @@ class ModuleHelper(object):
 
 class Module(object):
     def __init__(self,
+                 key=None,
                  url=None,
                  module_dir=None,
                  signatures=None,
@@ -105,7 +110,9 @@ class Module(object):
         self.processor = None
         self.name = "temp"
         # TODO(wuzewu): print more module loading info log
-        if url:
+        if key:
+            self._init_with_key(key=key)
+        elif url:
             self._init_with_url(url=url)
         elif module_dir:
             self._init_with_module_file(module_dir=module_dir)
@@ -124,10 +131,23 @@ class Module(object):
         else:
             raise "Error! HubModule can't init with nothing"
 
+    def _init_with_key(self, key):
+        logger.info("Try installing module %s" % key)
+        result, tips, module_dir = default_module_manager.install_module(
+            module_name=key)
+        if not result:
+            logger.error(tips)
+            exit(1)
+        logger.info(tips)
+        self._init_with_module_file(module_dir)
+
     def _init_with_url(self, url):
         utils.check_url(url)
-        result, _, module_dir = default_downloader.download_file_and_uncompress(
+        result, tips, module_dir = default_downloader.download_file_and_uncompress(
             url, save_path=".")
+        if not result:
+            logger.error(tips)
+            exit(1)
         self._init_with_module_file(module_dir)
 
     def _dump_processor(self):
-- 
GitLab