Automatically download models

LetsData-net · Dec 8, 2021 · 181b970 · 181b970
1 parent f546cc6
commit 181b970
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/DeepNewsSentiment/dataset.py b/DeepNewsSentiment/dataset.py
@@ -39,7 +39,11 @@
 from DeepNewsSentiment.models.FXBaseModel import FXBaseModel
 
 logger = get_logger()
-nlp = spacy.load("en_core_web_sm")
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    spacy.cli.download("en_core_web_sm")
+    nlp = spacy.load("en_core_web_sm")
 # this is what we could run on spacy 2.2+
 # nlp_dep_parser_labels = list(nlp.parser.labels)
 # but since we're on spacy 2.1 (and need to be because of newsalyze-backend)

diff --git a/DeepNewsSentiment/download.py b/DeepNewsSentiment/download.py
@@ -1,20 +1,15 @@
 """
 Download a specific version of a finetuned model and place it in pretrained_models.
 """
-import argparse
 import os
 import string
-from contextlib import suppress
 
 import torch
 
 from DeepNewsSentiment.fxlogger import get_logger
 
 
 class Download:
-    MODEL_DIRECTORY = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "pretrained_models", "state_dicts"
-    )
 
     def __init__(
         self, own_model_name, version="default", force=False, list_versions=False
@@ -76,7 +71,10 @@ def model_filename(model_cls, version=None):
 
     @classmethod
     def model_path(cls, model_cls, version=None):
-        return os.path.join(cls.MODEL_DIRECTORY, cls.model_filename(model_cls, version))
+        return os.path.join(torch.hub.get_dir(),
+                            "pretrained_models",
+                            "state_dicts",
+                            cls.model_filename(model_cls, version))
 
     @staticmethod
     def add_subparser(subparser):

diff --git a/DeepNewsSentiment/hubconf.py → hubconf.py b/DeepNewsSentiment/hubconf.py → hubconf.py