From 1e155602fadb463099ef16dd0ce26d86f11fd1ee Mon Sep 17 00:00:00 2001 From: markus583 Date: Tue, 10 Sep 2024 15:03:56 +0200 Subject: [PATCH] remove if __name__ --- wtpsplit/__init__.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/wtpsplit/__init__.py b/wtpsplit/__init__.py index cf643a14..538c2bfe 100644 --- 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!")