remove if __name__

segment-any-text · Sep 10, 2024 · 1e15560 · 1e15560
1 parent 88bdc91
commit 1e15560
Showing 1 changed file with 3 additions and 10 deletions.
diff --git a/wtpsplit/__init__.py b/wtpsplit/__init__.py
@@ -453,7 +453,9 @@ def __init__(
             else:
                 # no need to load if no ort_providers set
                 if ort_providers is not None:
-                    onnx_path = cached_file(model_name_to_fetch, "model_optimized.onnx", **(from_pretrained_kwargs or {}))
+                    onnx_path = cached_file(
+                        model_name_to_fetch, "model_optimized.onnx", **(from_pretrained_kwargs or {})
+                    )
                 else:
                     onnx_path = None
 
@@ -790,12 +792,3 @@ def get_default_threshold(model_str: str):
                     text, np.where(probs > sentence_threshold)[0], strip_whitespace=strip_whitespace
                 )
                 yield sentences
-
-
-if __name__ == "__main__":
-    sat = SaT("sat-3l-lora", ort_providers=["CPUExecutionProvider"])
-    print(sat.split("Hello, World! Next."))
-
-    wtp = WtP("wtp-bert-tiny", ort_providers=["CPUExecutionProvider"])
-    print(wtp.split("Hello, World! Next."))
-    print("DONE!")