pipeline classes

src/pipelines/fasttext_pipeline.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+
+import fasttext
+
+from src.models.fasttext import predict_samples
+
+
+class FastTextPipeline:
+    def __init__(self, model_path: Path, model_name: str):
+        self.model = fasttext.load_model(str(model_path))
+        self.model_name = model_name
+
+    def run(
+        self,
+        samples: list[dict],
+        max_len: int = 256,
+        batch_size: int = 32,
+        deduplicate: bool = False,
+    ) -> list[dict]:
+        return predict_samples(
+            self.model, samples, deduplicate=deduplicate,
+        )

src/pipelines/predict_all_pipeline.py

@@ -0,0 +1,57 @@
+import logging
+from pathlib import Path
+
+import onnxruntime as ort
+
+from src.models.inference import build_tokenizer, predict
+from src.pipelines.fasttext_pipeline import FastTextPipeline
+
+log = logging.getLogger(__name__)
+
+
+class PredictAllPipeline:
+    def __init__(self):
+        self.models: dict[str, dict] = {}
+        self.fasttext_models: dict[str, FastTextPipeline] = {}
+
+    def add_model(
+        self,
+        name: str,
+        onnx_path: Path,
+        mode: str,
+    ) -> None:
+        log.info(f"Adding {name} to PredictAllPipeline")
+        self.models[name] = {
+            "session": ort.InferenceSession(str(onnx_path)),
+            "tokenizer": build_tokenizer(mode),
+            "mode": mode,
+        }
+
+    def add_fasttext(self, name: str, model_path: Path) -> None:
+        log.info(f"Adding {name} (fastText) to PredictAllPipeline")
+        self.fasttext_models[name] = FastTextPipeline(model_path, name)
+
+    def run(
+        self,
+        samples: list[dict],
+        max_len: int = 256,
+        batch_size: int = 32,
+        deduplicate: bool = False,
+    ) -> dict[str, list[dict]]:
+        results: dict[str, list[dict]] = {}
+        for name, m in self.models.items():
+            results[name] = predict(
+                samples=samples,
+                session=m["session"],
+                tokenizer=m["tokenizer"],
+                mode=m["mode"],
+                max_len=max_len,
+                batch_size=batch_size,
+                deduplicate=deduplicate,
+            )
+        for name, ft in self.fasttext_models.items():
+            results[name] = ft.run(
+                samples=samples,
+                deduplicate=deduplicate,
+            )
+        return results

src/pipelines/predict_pipeline.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+
+import onnxruntime as ort
+
+from src.models.inference import build_tokenizer, predict
+
+
+class PredictPipeline:
+    def __init__(
+        self,
+        onnx_path: Path,
+        mode: str,
+        model_name: str,
+    ):
+        self.session = ort.InferenceSession(str(onnx_path))
+        self.tokenizer = build_tokenizer(mode)
+        self.mode = mode
+        self.model_name = model_name
+
+    def run(
+        self,
+        samples: list[dict],
+        max_len: int = 256,
+        batch_size: int = 32,
+        deduplicate: bool = False,
+    ) -> list[dict]:
+        return predict(
+            samples=samples,
+            session=self.session,
+            tokenizer=self.tokenizer,
+            mode=self.mode,
+            max_len=max_len,
+            batch_size=batch_size,
+            deduplicate=deduplicate,
+        )