✨ English model

Author: pajowu

generate_models.py

@@ -9,6 +9,8 @@
 import pickle
 
 from punctuator import Punctuator
+import gdown
+
 
 CACHE_PATH = Path.home() / "cache"
 CACHE_PATH.mkdir(exist_ok=True)
@@ -19,22 +21,63 @@
     {
         "url": "http://ltdata1.informatik.uni-hamburg.de/subtitle2go/Model_subs_norm1_filt_5M_tageschau_euparl_h256_lr0.02.pcl",
         "tests": [{"input": "hallo ich bin ein testsatz", "expected": "Hallo, ich bin ein testsatz."}]
+    },
+    # Rehosted from https://drive.google.com/drive/folders/0B7BsN5f2F1fZQnFsbzJ3TWxxMms?resourcekey=0-6yhuY9FOeITBBWWNdyG2aw
+    {
+        "url": "gdrive://1CZ_Os38LjBwyd-jgDMsfpqiWPB6wwVKA",
+        "name": "Demo-EUROPARL-EN.zip",
+        "pickle_encoding": "latin-1",
+        "tests": [
+            {
+                "input": "hello this is an example sentence",
+                "expected": "Hello, this is an example sentence.",
+            }
+        ],
     }
 ]
 
-for model in PUNCTUATOR_MODELS:
-    url_path = urllib.parse.urlparse(model['url']).path
+
+def download_model(model):
+    if model["url"].startswith("gdrive://"):
+        return download_gdrive_model(model)
+    else:
+        return download_http_model(model)
+
+
+def download_gdrive_model(model):
+    url_path = urllib.parse.urlparse(model["url"]).netloc
+    output_model_file_path = MODEL_PATH / model["name"]
+    input_model_file_path = CACHE_PATH / url_path
+    if not input_model_file_path.exists():
+        print("Downloading", url_path)
+        gdown.download(id=url_path, output=str(input_model_file_path), fuzzy=True)
+    return input_model_file_path, output_model_file_path
+
+
+def download_http_model(model):
+    url_path = urllib.parse.urlparse(model["url"]).path
     name = Path(url_path).name
     input_model_file_path = CACHE_PATH / name
     output_name = Path(url_path).with_suffix(".zip").name
     output_model_file_path = MODEL_PATH / output_name
     if not input_model_file_path.exists():
-        req = requests.get(model['url'])
+        print("Downloading", model["url"])
+        req = requests.get(model["url"])
         with open(input_model_file_path, "wb") as f:
             f.write(req.content)
 
+    return input_model_file_path, output_model_file_path
+
+
+for model in PUNCTUATOR_MODELS:
+    input_model_file_path, output_model_file_path = download_model(model)
     with open(input_model_file_path, "rb") as f:
-        state = pickle.load(f)
+        if 'pickle_encoding' in model:
+            u = pickle._Unpickler(f)
+            u.encoding = model['pickle_encoding']
+            state = u.load()
+        else:
+            state = pickle.load(f)
 
     with zipfile.ZipFile(output_model_file_path, "w") as model_zip:
         for k, v in state.items():
@@ -50,8 +93,10 @@
                 with model_zip.open(f"{k}.json", "w") as f:
                     f.write(json.dumps(v).encode())
 
-    if 'tests' in model and model['tests']:
+    if "tests" in model and model["tests"]:
         punctuation_model = Punctuator(output_model_file_path)
-        for test in model['tests']:
-            actual = punctuation_model.punctuate(test['input'])
-            assert actual == test['expected']
+        for test in model["tests"]:
+            actual = punctuation_model.punctuate(test["input"])
+            assert (
+                actual == test["expected"]
+            ), f"'{test['expected']}' expected, got {actual}"

poetry.lock

@@ -10,6 +10,7 @@ numpy = ">1.22"
 scipy = ">=1.7.0"
 punctuator = {git = "https://github.com/audapolis/punctuator2", rev = "51eeb14cd30a4162e6a54e5234805651963f767f"}
 requests = "^2.27.1"
+gdown = "^4.4.0"
 
 [tool.poetry.dev-dependencies]