minor adaptions

Author: Johannes Hötter

refinery/callbacks/__init__.py

@@ -0,0 +1,40 @@
+import pandas as pd
+from refinery import exceptions
+
+
+class ModelCallback:
+    def __init__(
+        self, client, inference_fn, preprocessing_fn=None, postprocessing_fn=None
+    ):
+        self.client = client
+        self.inference_fn = inference_fn
+        self.preprocessing_fn = preprocessing_fn
+        self.postprocessing_fn = postprocessing_fn
+
+        self.primary_keys = client.get_primary_keys()
+
+    @staticmethod
+    def __batch(documents):
+        BATCH_SIZE = 32
+        length = len(documents)
+        for idx in range(0, length, BATCH_SIZE):
+            yield documents[idx : min(idx + BATCH_SIZE, length)]
+
+    def run(self, inputs, indices):
+        indices_df = pd.DataFrame(indices)
+        if not all([key in indices_df.columns for key in self.primary_keys]):
+            raise exceptions.PrimaryKeyError("Errorneous primary keys given for index.")
+
+        index_generator = ModelCallback.__batch(indices)
+        for batched_inputs in ModelCallback.__batch(inputs):
+            batched_indices = next(index_generator)
+
+            if self.preprocessing_fn is not None:
+                batched_inputs = self.preprocessing_fn(batched_inputs)
+
+            batched_outputs = self.inference_fn(batched_inputs)
+
+            if self.postprocessing_fn is not None:
+                batched_outputs = self.postprocessing_fn(batched_outputs)
+
+            yield {"index": batched_indices, "associations": batched_outputs}

refinery/callbacks/inference.py

@@ -1,12 +1,19 @@
 # -*- coding: utf-8 -*-
 from typing import Optional
 
+
 class LocalError(Exception):
     pass
 
+
 class UnknownItemError(LocalError):
     pass
 
+
+class PrimaryKeyError(LocalError):
+    pass
+
+
 # https://developer.mozilla.org/en-US/docs/Web/HTTP/Status#client_error_responses
 class APIError(Exception):
     def __init__(self, message: Optional[str] = None):
@@ -24,9 +31,12 @@ class UnauthorizedError(APIError):
 class NotFoundError(APIError):
     pass
 
+
 class UnknownProjectError(APIError):
     def __init__(self, project_id: str):
-        super().__init__(message=f"Could not find project '{project_id}'. Please check your input.")
+        super().__init__(
+            message=f"Could not find project '{project_id}'. Please check your input."
+        )
 
 
 # 500 Server Error
@@ -37,9 +47,10 @@ class InternalServerError(APIError):
 class FileImportError(Exception):
     pass
 
+
 # mirror this from the rest api class ErrorCodes
 class ErrorCodes:
-    UNRECOGNIZED_USER = "UNRECOGNIZED_USER" # not actively used in SDK
+    UNRECOGNIZED_USER = "UNRECOGNIZED_USER"  # not actively used in SDK
     PROJECT_NOT_FOUND = "PROJECT_NOT_FOUND"