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

@@ -1,27 +1,30 @@
 import re
 from typing import Dict, Iterable, List
 
-from helper import preprocess
+from config import model_key
+from preprocess import preprocess
 from models import DomainPrediction
 from sklearn.pipeline import Pipeline
 
-# from sus.use_model import use_model
+from good_ai import use_model
+from sus.clean import clean
 
-
-# @use_model(model_key, version="latest")
-def predict(
+@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 re.sub(r"[^a-zA-Z0-9]", " ", text).split(" ")
+        for original in text.split(" ")
     }
 
-    features = model.named_steps["vectorizer"].transform([text])
+    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)
 

example/preprocess.py

@@ -5,6 +5,5 @@ from sus.lemmatize_text import lemmatize_text
 
 
 def preprocess(text: str) -> str:
-    cleaned = clean(text, convert_to_ascii=True)
-    lemmas = [re.sub(r"\d+", "NUM", lemma) for lemma in lemmatize_text(cleaned)]
+    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/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)
 

good_ai/src/sus.egg-info/PKG-INFO

@@ -1,50 +0,0 @@
-Metadata-Version: 2.1
-Name: sus
-Version: 0.0.3
-Summary: [S]coutinScience [u]tilitie[s]: reusable utilities for text processing
-Home-page: https://github.com/ScoutinScience/platform
-Author: ScoutinScience B.V.
-Author-email: andras@scoutinscience.com
-License: UNKNOWN
-Project-URL: Bug Tracker, https://github.com/ScoutinScience/platform/issues
-Platform: UNKNOWN
-Classifier: Programming Language :: Python :: 3
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-
-# **S**coutinScience **U**tilitie**S** for text processing [![Lint and test ScoutinScience utilities](https://github.com/ScoutinScience/platform/actions/workflows/sus-general.yaml/badge.svg)](https://github.com/ScoutinScience/platform/actions/workflows/sus-general.yaml)
-
-> amogus
-
-## Exports
-
-- [clean](src/sus/clean.py)
-- [unique](src/sus/unique.py)
-- [parallel_map](src/sus/parallel_map.py)
-- [match_names](src/sus/match_names/match_names.py)
-- [evaluate_ranking](src/sus/evaluate_ranking/evaluate_ranking.py)
-
-### Requires loading spacy model
-
-> This is automatic but will require some time.
-
-> Add this to the Dockerfile for caching the spaCy model:
->
-> ```docker
-> RUN pip install --no-cache-dir en-core-web-lg@https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.2.0/en_core_web_lg-3.2.0-py3-none-any.whl
-> ```
-
-- [spacy model (nlp)](src/sus/nlp.py)
-- [get_sentences](src/sus/get_sentences.py)
-- [lemmatize_text](src/sus/lemmatize_text.py)
-- [lemmatize_token](src/sus/lemmatize_token.py)
-- [publication TEI](src/sus/publication_tei/publication_tei.py)
-
-## Development
-
-- Optional booleans must have a default value of `False`.
-- No imports in top-level `__init__.py`, in order to not load anything unnecessary automatically
-- Should only be updated through a PR
-
-

good_ai/src/sus.egg-info/SOURCES.txt

@@ -1,78 +0,0 @@
-.gitignore
-README.md
-pyproject.toml
-requirements.txt
-setup.cfg
-src/__init__.py
-src/sus/__init__.py
-src/sus/clean.py
-src/sus/get_sentences.py
-src/sus/lemmatize_text.py
-src/sus/lemmatize_token.py
-src/sus/nlp.py
-src/sus/parallel_map.py
-src/sus/unique.py
-src/sus.egg-info/PKG-INFO
-src/sus.egg-info/SOURCES.txt
-src/sus.egg-info/dependency_links.txt
-src/sus.egg-info/requires.txt
-src/sus.egg-info/top_level.txt
-src/sus/data/__init__.py
-src/sus/data/american_spellings.py
-src/sus/data/punctuations.py
-src/sus/evaluate_ranking/__init__.py
-src/sus/evaluate_ranking/draw_f1_iso_lines.py
-src/sus/evaluate_ranking/evaluate_ranking.py
-src/sus/external/__init__.py
-src/sus/external/negspacy/README.md
-src/sus/external/negspacy/__init__.py
-src/sus/external/negspacy/negation.py
-src/sus/external/negspacy/termsets.py
-src/sus/external/pylatexenc/README.md
-src/sus/external/pylatexenc/__init__.py
-src/sus/external/pylatexenc/_util.py
-src/sus/external/pylatexenc/version.py
-src/sus/external/pylatexenc/latex2text/__init__.py
-src/sus/external/pylatexenc/latex2text/__main__.py
-src/sus/external/pylatexenc/latex2text/_defaultspecs.py
-src/sus/external/pylatexenc/latexencode/__init__.py
-src/sus/external/pylatexenc/latexencode/__main__.py
-src/sus/external/pylatexenc/latexencode/_partial_latex_encoder.py
-src/sus/external/pylatexenc/latexencode/_uni2latexmap.py
-src/sus/external/pylatexenc/latexencode/_uni2latexmap_xml.py
-src/sus/external/pylatexenc/latexencode/_unicode_to_latex_encoder.py
-src/sus/external/pylatexenc/latexwalker/__init__.py
-src/sus/external/pylatexenc/latexwalker/__main__.py
-src/sus/external/pylatexenc/latexwalker/_defaultspecs.py
-src/sus/external/pylatexenc/macrospec/__init__.py
-src/sus/external/pylatexenc/macrospec/_argparsers.py
-src/sus/language/__init__.py
-src/sus/language/english_name_of_language.py
-src/sus/language/is_english.py
-src/sus/language/predict_language.py
-src/sus/match_names/__init__.py
-src/sus/match_names/config.py
-src/sus/match_names/match_names.py
-src/sus/match_names/name_parts.py
-src/sus/publication_tei/__init__.py
-src/sus/publication_tei/publication_tei.py
-src/sus/publication_tei/models/__init__.py
-src/sus/publication_tei/models/affiliation.py
-src/sus/publication_tei/models/author.py
-src/sus/publication_tei/models/element.py
-src/sus/publication_tei/models/publication_metadata.py
-src/sus/publication_tei/models/text.py
-tests/__init__.py
-tests/test_clean.py
-tests/test_evaluate_ranking.py
-tests/test_get_sentences.py
-tests/test_language.py
-tests/test_lemmatize_text.py
-tests/test_lemmatize_token.py
-tests/test_match_names.py
-tests/test_parallel_map.py
-tests/test_publication_tei.py
-tests/test_unique.py
-tests/data/10.1136_bmjspcare-2021-003026.pdf.tei.xml
-tests/data/bad.tei.xml
-tests/data/parsed.py

good_ai/src/sus.egg-info/requires.txt

@@ -1,14 +0,0 @@
-click<8.1.0
-unidecode>=1.3.0
-multiprocess>=0.70.0.0
-tqdm>=4.0.0
-psutil>=5.9.0
-beautifulsoup4>=4.10.0
-lxml>=4.6.0
-spacy>=3.2.0
-pydantic>=1.8.0
-scikit-learn>=1.0.0
-matplotlib>=3.5.0
-numpy>=1.22.0
-langcodes[data]>=3.3.0
-langdetect>=1.0.9

good_ai/src/sus.egg-info/top_level.txt

@@ -1 +0,0 @@
-sus

good_ai/src/sus/parallel_map.py

@@ -10,9 +10,11 @@ def parallel_map(
     function: Callable[[Any], Any],
     values: Iterable[Any],
     chunk_size: Optional[int] = None,
-    concurrency: int = psutil.cpu_count(),
+    concurrency: Optional[int] = None,
     disable_progress: bool = False,
 ) -> List[Any]:
+    if concurrency is None:
+        concurrency = psutil.cpu_count()
     assert concurrency > 0
     assert chunk_size is None or chunk_size > 0