Experiment with framework

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

.gitignore

@@ -2,3 +2,5 @@
 .DS_Store
 __pycache__
 .cache
+.mypy_cache
+.pytest_cache

example/helper/__init__.py

@@ -1 +0,0 @@
-from .preprocess import preprocess

example/main.py

@@ -0,0 +1,16 @@
+import json
+
+from open_s3 import LargeFile
+from predict import predict_domain
+
+from good_ai import process_batch
+
+if __name__ == "__main__":
+    with open("data/s2-corpus-1583.json") as f:
+        raw = json.load(f)
+
+    LargeFile.configure_credentials_from_file("s3.ini")
+
+    data = {f'{r["title"]} {r["abstract"]}': r["domain"] for r in raw[:5]}
+
+    print(process_batch(predict_domain, ["We have found a new type of chemical."]))

example/predict_domain.py

@@ -0,0 +1,68 @@
+import re
+from typing import Dict, Iterable, List
+
+from config import model_key
+from preprocess import preprocess
+from models import DomainPrediction
+from sklearn.pipeline import Pipeline
+
+from good_ai import use_model
+from sus.clean import clean
+
+@use_model(model_key, version="latest")
+def predict_domain(
+    text: str, model: Pipeline, cut_off_probability: float = 0.2
+) -> List[DomainPrediction]:
+    assert 0 <= cut_off_probability <= 1
+    cleaned = clean(text, convert_to_ascii=True)
+    text = re.sub(r"[^a-zA-Z0-9]", " ", cleaned)
+
+    feature_names = model.named_steps["vectorizer"].get_feature_names_out()
+
+    token_mapping = {
+        preprocess(original): original
+        for original in text.split(" ")
+    }
+
+    features = model.named_steps["vectorizer"].transform([" ".join(token_mapping.keys())])
+    prediction = model.named_steps["classifier"].predict_proba(features)[0]
+    best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
+
+    results: List[DomainPrediction] = []
+    for class_index, probability in best_classes:
+        weights = model.named_steps["classifier"].feature_log_prob_[class_index]
+
+        results.append(
+            DomainPrediction(
+                domain=model.named_steps["classifier"].classes_[class_index],
+                probability=round(probability * 100),
+                explanation=_get_explanation(
+                    feature_names=feature_names,
+                    features=features.A[0],
+                    weights=weights,
+                    token_mapping=token_mapping,
+                ),
+            )
+        )
+
+        if sum(r.probability for r in results) >= cut_off_probability:
+            break
+
+    return results
+
+
+def _get_explanation(
+    feature_names: Iterable[str],
+    features: Iterable[float],
+    weights: Iterable[float],
+    token_mapping: Dict[str, str],
+) -> List[str]:
+    influential = [
+        (value * weight, name)
+        for name, value, weight in zip(feature_names, features, weights)
+        if value
+    ]
+
+    most_influential = sorted(influential, reverse=True)[:5]
+
+    return [token_mapping[v[1]] for v in most_influential]

example/preprocess.py

@@ -0,0 +1,9 @@
+import re
+
+from sus.clean import clean
+from sus.lemmatize_text import lemmatize_text
+
+
+def preprocess(text: str) -> str:
+    lemmas = [re.sub(r"\d+", "NUM", lemma) for lemma in lemmatize_text(text)]
+    return " ".join(lemmas)

example/secrets-backblaze.ini

@@ -1,6 +0,0 @@
-[DEFAULT]
-aws_region_name = us-west-4
-aws_access_key_id = 00411a2454be9f90000000001
-aws_secret_access_key = K004YOSYHePVUHymVzm/vb+AjGahXl8
-large_files_bucket_name = sa-large-files
-endpoint_url = https://s3.us-west-004.backblazeb2.com

good_ai/src/good_ai/__init__.py

@@ -0,0 +1,2 @@
+from .deploy import process_batch
+from .models import save_model, use_model

good_ai/src/good_ai/core/__init__.py

@@ -0,0 +1,2 @@
+from .function_registry import function_registry
+from .plugin import Plugin

good_ai/src/good_ai/core/function_registry.py

@@ -0,0 +1,23 @@
+from collections import defaultdict
+from typing import Any, Callable, DefaultDict, List
+
+from .plugin import Plugin
+
+
+class FunctionRegistry:
+    def __init__(self) -> None:
+        self._registered_functions: DefaultDict[int, List[Plugin]] = defaultdict(
+            lambda: []
+        )
+
+    def add_plugin(self, function: Callable[..., Any], plugin: Plugin):
+        self._registered_functions[id(function)].append(plugin)
+
+    def get_plugins(self, function: Callable[..., Any]) -> List[Plugin]:
+        plugins = self._registered_functions[id(function)]
+        for p in plugins:
+            p.initialize()
+        return plugins
+
+
+function_registry = FunctionRegistry()

good_ai/src/good_ai/core/plugin.py

@@ -0,0 +1,19 @@
+from typing import Any, Callable
+
+
+class Plugin:
+    def __init__(self, function: Callable[[Any], Any]):
+        self._function = function
+        self._initialized = False
+
+    def initialize(self):
+        if not self._initialized:
+            self.on_initialize()
+            self._initialized = True
+
+    def on_initialize(self):
+        pass
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        assert self._initialized
+        return self._function(*args, **kwargs)

good_ai/src/good_ai/deploy/__init__.py

@@ -0,0 +1 @@
+from .process_batch import process_batch

good_ai/src/good_ai/deploy/process_batch.py

@@ -0,0 +1,16 @@
+from functools import partial, reduce
+from typing import Any, Callable, Iterable, Optional, Sequence
+
+from sus.parallel_map import parallel_map
+
+from ..core import function_registry
+
+
+def process_batch(
+    function: Callable[..., Any],
+    batch: Iterable[Any],
+    concurrency: Optional[int] = None,
+) -> Sequence[Any]:
+    plugins = function_registry.get_plugins(function)
+    composed = partial(reduce, lambda r, f: f(r), plugins)
+    return parallel_map(composed, batch, concurrency=concurrency)

good_ai/src/good_ai/deploy/serve.py

@@ -0,0 +1 @@
+

good_ai/src/good_ai/models/__init__.py

@@ -0,0 +1,2 @@
+from .save_model import save_model
+from .use_model import use_model

good_ai/src/good_ai/models/load_model.py

@@ -0,0 +1,19 @@
+import logging
+from typing import Any, Optional
+
+from joblib import load
+from open_s3 import LargeFile
+
+logger = logging.getLogger("models")
+
+
+def load_model(
+    key: str, version: Optional[int] = None, return_path: bool = False
+) -> Any:
+    file = LargeFile(name=key, mode="rb", version=version)
+
+    if return_path:
+        return file.get()
+
+    with file as f:
+        return load(f)

good_ai/src/good_ai/models/save_model.py

@@ -0,0 +1,24 @@
+import logging
+from pathlib import Path
+from typing import Optional, Union
+
+from joblib import dump
+from open_s3 import LargeFile
+
+logger = logging.getLogger("models")
+
+
+def save_model(
+    model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
+) -> int:
+    file = LargeFile(name=key, mode="wb", keep_last_n=keep_last_n)
+
+    if isinstance(model, Path) or isinstance(model, str):
+        file.push(model)
+    else:
+        with file as f:
+            dump(model, f)
+
+    logger.info(f"Model {key} uploaded with version {file.version}")
+
+    return file.version

good_ai/src/good_ai/models/use_model.py

@@ -0,0 +1,24 @@
+from typing import Literal, Union
+
+from ..core import function_registry
+from .use_model_plugin import UseModelPlugin
+
+
+def use_model(
+    key: str, version: Union[int, Literal["latest"]] = None, return_path: bool = False
+):
+    assert isinstance(version, int) or version == "latest"
+
+    def inner(f):
+        function_registry.add_plugin(
+            f,
+            UseModelPlugin(
+                f,
+                key=key,
+                version=version if isinstance(version, int) else None,
+                return_path=return_path,
+            ),
+        )
+        return f
+
+    return inner

good_ai/src/good_ai/models/use_model_plugin.py

@@ -0,0 +1,28 @@
+from typing import Any, Callable, Optional
+
+from ..core import Plugin
+from .load_model import load_model
+
+
+class UseModelPlugin(Plugin):
+    def __init__(
+        self,
+        function: Callable[[Any], Any],
+        key: str,
+        version: Optional[int],
+        return_path: bool,
+    ):
+        def wrapper(*args, **kwargs):
+            return function(*args, **kwargs, model=self._model)
+
+        super().__init__(wrapper)
+
+        self._key = key
+        self._version = version
+        self._return_path = return_path
+
+    def on_initialize(self) -> None:
+        super().on_initialize()
+        self._model = load_model(
+            key=self._key, version=self._version, return_path=self._return_path
+        )

good_ai/src/open_s3/large_file.py

@@ -8,7 +8,8 @@ from types import TracebackType
 from typing import IO, Any, Dict, List, Optional, Type, Union
 
 import boto3
-from helper import DownloadProgressBar, UploadProgressBar, human_readable_to_byte
+
+from .helper import DownloadProgressBar, UploadProgressBar, human_readable_to_byte
 
 logger = logging.getLogger("open_s3")
 
@@ -158,6 +159,10 @@ class LargeFile:
     def version_ids(self) -> List[int]:
         return [self._get_version_from_key(key) for key in self._versions]
 
+    @property
+    def version(self) -> int:
+        return self._version
+
     def get(self, hide_progress: bool = False) -> Path:
         key = next(
             key
@@ -194,7 +199,7 @@ class LargeFile:
 
         return destination
 
-    def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
+    def push(self, path: Union[Path, str], hide_progress: bool = False) -> int:
         if isinstance(path, str):
             path = Path(path)