diff --git a/mailcom/parse.py b/mailcom/parse.py index 75171dd..81ea972 100644 --- a/mailcom/parse.py +++ b/mailcom/parse.py @@ -67,9 +67,21 @@ def init_spacy(self, language: str, model="default"): model, exclude=["morphologizer", "attribute_ruler", "lemmatizer", "ner"] ) except OSError: - raise OSError("Could not find {} in standard directory.".format(model)) - - self.nlp_spacy = sp.load(model) + pass + try: + print( + "Could not find model in standard directory. Trying to download model from repo." # noqa + ) + # try downloading model + sp.cli.download(model) + self.nlp_spacy = sp.load( + model, + exclude=["morphologizer", "attribute_ruler", "lemmatizer", "ner"], + ) + except SystemExit: + raise SystemExit("Could not download {} from repo".format(model)) + except OSError: + raise OSError("Could not find {} in standard directory".format(model)) def init_transformers( self, diff --git a/mailcom/test/test_parse.py b/mailcom/test/test_parse.py index 6383880..b7a782e 100644 --- a/mailcom/test/test_parse.py +++ b/mailcom/test/test_parse.py @@ -44,7 +44,7 @@ def get_default_fr(): def test_init_spacy(get_instant): with pytest.raises(KeyError): get_instant.init_spacy("not_a_language") - with pytest.raises(OSError): + with pytest.raises(SystemExit): get_instant.init_spacy("fr", "not_an_existing_spacy_model")