From 8f23d396b9240213fbecdc22f4f7cac985d890d5 Mon Sep 17 00:00:00 2001 From: Piotr Czapla Date: Mon, 11 Nov 2019 09:51:11 +0100 Subject: [PATCH] Make form_pretrained accept a path as well as a pretarined model name --- multifit/training.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/multifit/training.py b/multifit/training.py index c2bd4ff..cb4c9c9 100644 --- a/multifit/training.py +++ b/multifit/training.py @@ -593,9 +593,12 @@ class ULMFiT: return results def from_pretrained_(self, name, repo="n-waves/multifit-models"): - name = name.rstrip(".tgz") # incase someone put's tgz name the name - url = f"https://github.com/{repo}/releases/download/{name}/{name}.tgz" - path = untar_data(url.rstrip(".tgz"), data=False) # untar_data adds .tgz + if (Path(name)/f"{LM_BEST}.pth").exists(): + path = Path(name) + else: + name = name.rstrip(".tgz") # incase someone put's tgz name the name + url = f"https://github.com/{repo}/releases/download/{name}/{name}.tgz" + path = untar_data(url.rstrip(".tgz"), data=False) # untar_data adds .tgz return self.load_(path)