More experimentation

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

.vscode/settings.json

@@ -1,3 +1,19 @@
 {
-  "cSpell.words": ["pydantic", "pyplot", "sklearn", "Tfidf", "Vectorizer"]
+  "cSpell.words": [
+    "botocore",
+    "pydantic",
+    "pyplot",
+    "sklearn",
+    "Tfidf",
+    "tinydb",
+    "Vectorizer",
+    "xmargin"
+  ],
+  "files.exclude": {
+    ".env": true,
+    "**/.cache": true,
+    "**/.ipynb_checkpoints": true,
+    "**/.mypy_cache": true,
+    "**/.pytest_cache": true
+  }
 }

example/.gitignore

@@ -1,2 +1,3 @@
 data
 .ipynb_checkpoints
+tracing_database.json

example/README.md

@@ -0,0 +1,12 @@
+# Train Domain classifier from the [semantic scholar dataset](https://api.semanticscholar.org/corpus)
+
+## Upload the dataset (or a part of it) to shared infrastructure
+
+```sh
+mkdir ss-data && cd ss-data
+wget https://s3-us-west-2.amazonaws.com/ai2-s2-research-public/open-corpus/2022-02-01/manifest.txt
+wget -B https://s3-us-west-2.amazonaws.com/ai2-s2-research-public/open-corpus/2022-02-01/ -i manifest.txt
+cd -
+python3 -m good_ai.open_s3 --secrets s3.ini --push ss-data
+rm -rf ss-data
+```

example/main.py

@@ -1,15 +1,21 @@
 import json
+from random import shuffle
 
+from devtools import debug
 from predict_domain import predict_domain
 
-from good_ai import LargeFile, process_batch
+from good_ai import process_batch
 
 if __name__ == "__main__":
-    with open("data/s2-corpus-1583.json") as f:
+    with open(".cache/ss-data-0/s2-corpus-1583.json") as f:
         raw = json.load(f)
 
-    LargeFile.configure_credentials_from_file("s3.ini")
+    shuffle(raw)
-
     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."]))
+    results = process_batch(predict_domain, data.keys())
+
+    for predicted, actual in zip(results, data.values()):
+        print(", ".join(actual))
+        debug(predicted)
+        print()

example/predict_domain.py

@@ -45,7 +45,7 @@ def predict_domain(
             )
         )
 
-        if sum(r.probability for r in results) >= cut_off_probability:
+        if sum(r.probability for r in results) >= cut_off_probability * 100:
             break
 
     return results

good_ai/src/good_ai/good_ai/__init__.py

@@ -1,2 +1,3 @@
 from .deploy import process_batch
 from .models import save_model, use_model
+from .set_default_config import set_default_config, set_default_config_if_uninitialized

good_ai/src/good_ai/good_ai/core/__init__.py

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

good_ai/src/good_ai/good_ai/core/function_registry.py

@@ -1,23 +0,0 @@
-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/good_ai/core/plugin.py

@@ -1,19 +0,0 @@
-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/good_ai/deploy/process_batch.py

@@ -1,16 +1,23 @@
-from functools import partial, reduce
 from typing import Any, Callable, Iterable, Optional, Sequence
 
 from good_ai.utilities.parallel_map import parallel_map
 
-from ..core import function_registry
+from ..set_default_config import set_default_config_if_uninitialized
+from ..tracing import Trace, TracingContext
 
 
 def process_batch(
     function: Callable[..., Any],
     batch: Iterable[Any],
     concurrency: Optional[int] = None,
-) -> Sequence[Any]:
+) -> Sequence[Trace]:
-    plugins = function_registry.get_plugins(function)
+    set_default_config_if_uninitialized()
-    composed = partial(reduce, lambda r, f: f(r), plugins)
+
-    return parallel_map(composed, batch, concurrency=concurrency)
+    def inner(input: Any) -> Trace:
+        with TracingContext() as t:
+            t.log_input(input)
+            result = function(input)
+            output = t.log_output(result)
+        return output
+
+    return parallel_map(inner, batch, concurrency=concurrency)

good_ai/src/good_ai/good_ai/models/load_model.py

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

good_ai/src/good_ai/good_ai/models/save_model.py

@@ -6,12 +6,16 @@ from joblib import dump
 
 from good_ai.open_s3 import LargeFile
 
+from ..set_default_config import set_default_config_if_uninitialized
+
 logger = logging.getLogger("models")
 
 
 def save_model(
     model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
 ) -> int:
+    set_default_config_if_uninitialized()
+
     file = LargeFile(name=key, mode="wb", keep_last_n=keep_last_n)
 
     if isinstance(model, Path) or isinstance(model, str):

good_ai/src/good_ai/good_ai/models/use_model.py

@@ -1,24 +1,33 @@
-from typing import Literal, Union
+from functools import wraps
+from typing import Any, Callable, Literal, Union
 
-from ..core import function_registry
+from ..tracing import Model, TracingContext
-from .use_model_plugin import UseModelPlugin
+from .load_model import load_model
 
 
 def use_model(
-    key: str, version: Union[int, Literal["latest"]] = None, return_path: bool = False
+    key: str,
-):
+    *,
+    version: Union[int, Literal["latest"]],
+    return_path: bool = False,
+    model_kwarg_name: str = "model"
+) -> Callable[..., Any]:
     assert isinstance(version, int) or version == "latest"
 
-    def inner(f):
+    model, actual_version = load_model(
-        function_registry.add_plugin(
+        key=key,
-            f,
+        version=None if version == "latest" else version,
-            UseModelPlugin(
+        return_path=return_path,
-                f,
+    )
-                key=key,
-                version=version if isinstance(version, int) else None,
-                return_path=return_path,
-            ),
-        )
-        return f
 
-    return inner
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            context = TracingContext.get_current_context()
+            if context:
+                context.log_model(Model(key=key, version=actual_version))
+            return func(*args, **kwargs, **{model_kwarg_name: model})
+
+        return wrapper
+
+    return decorator

good_ai/src/good_ai/good_ai/models/use_model_plugin.py

@@ -1,28 +0,0 @@
-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/good_ai/good_ai/set_default_config.py

@@ -0,0 +1,55 @@
+import logging
+import random
+from pathlib import Path
+
+from good_ai.good_ai.tracing.tracing_context import TracingContext
+from good_ai.open_s3 import LargeFile
+
+from .tracing import PersistenceDriver, TinyDbDriver
+
+logger = logging.getLogger("good_ai")
+
+_initialized = False
+
+
+def set_default_config_if_uninitialized() -> None:
+    if not _initialized:
+        set_default_config()
+
+
+def set_default_config(
+    log_level: int = logging.INFO,
+    s3_config: Path = Path("s3.ini"),
+    seed: int = 42,
+    tracing_db_driver: PersistenceDriver = TinyDbDriver(Path("tracing_database.json")),
+) -> None:
+    global _initialized
+    logging.basicConfig(level=log_level)
+
+    _initialize_large_file(s3_config)
+    _set_seed(seed)
+
+    TracingContext.persistence_driver = tracing_db_driver
+
+    _initialized = True
+
+    logger.info(f"Defaults: configured ✅")
+
+
+def _initialize_large_file(s3_config: Path) -> None:
+    if s3_config.exists():
+        LargeFile.configure_credentials_from_file(s3_config)
+    else:
+        logger.info(
+            f"Provided S3 config ({s3_config.resolve()}) not found, skipping LargeFile initialisation"
+        )
+
+
+def _set_seed(seed: int) -> None:
+    random.seed(seed)
+    try:
+        import numpy
+
+        numpy.random.seed(seed + 1)
+    except ImportError:
+        pass

good_ai/src/good_ai/good_ai/tracing/__init__.py

@@ -0,0 +1,4 @@
+from .model import Model
+from .persistence import MongoDbDriver, PersistenceDriver, TinyDbDriver
+from .trace import Trace
+from .tracing_context import TracingContext

good_ai/src/good_ai/good_ai/tracing/model.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    key: str
+    version: int

good_ai/src/good_ai/good_ai/tracing/persistence/__init__.py

@@ -0,0 +1,3 @@
+from .mongodb_driver import MongoDbDriver
+from .persistence_driver import PersistenceDriver
+from .tinydb_driver import TinyDbDriver

good_ai/src/good_ai/good_ai/tracing/persistence/mongodb_driver.py

@@ -0,0 +1,5 @@
+from .persistence_driver import PersistenceDriver
+
+
+class MongoDbDriver(PersistenceDriver):
+    pass

good_ai/src/good_ai/good_ai/tracing/persistence/persistence_driver.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class PersistenceDriver(ABC):
+    @abstractmethod
+    def save_document(self, document: Dict[str, Any]) -> str:
+        pass

good_ai/src/good_ai/good_ai/tracing/persistence/tinydb_driver.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+from typing import Any, Dict
+
+from tinydb import TinyDB
+
+from .persistence_driver import PersistenceDriver
+
+
+class TinyDbDriver(PersistenceDriver):
+    def __init__(self, path_to_db: Path) -> None:
+        super().__init__()
+        self._db = TinyDB(path_to_db)
+
+    def save_document(self, document: Dict[str, Any]) -> str:
+        return self._db.insert(document)

good_ai/src/good_ai/good_ai/tracing/trace.py

@@ -0,0 +1,14 @@
+from typing import Any, List
+
+from pydantic import BaseModel
+
+from .model import Model
+
+
+class Trace(BaseModel):
+    created: str
+    execution_time_ms: float
+    input: Any
+    models: List[Model]
+    output: Any
+    evaluation: Any = None

good_ai/src/good_ai/good_ai/tracing/tracing_context.py

@@ -0,0 +1,71 @@
+import logging
+import threading
+from collections import defaultdict
+from datetime import datetime
+from types import TracebackType
+from typing import Any, DefaultDict, List, Optional, Type
+
+from .persistence import PersistenceDriver
+
+logger = logging.getLogger("good_ai")
+
+from .model import Model
+from .trace import Trace
+
+
+class TracingContext:
+    persistence_driver: PersistenceDriver
+    _contexts: DefaultDict[int, List["TracingContext"]] = defaultdict(lambda: [])
+
+    def __init__(self) -> None:
+        self._models: List[Model] = []
+        self._input: Any = None
+        self._output: Any = None
+        self._trace: Optional[Trace] = None
+        self._start_time = datetime.utcnow()
+
+    def log_input(self, input: Any) -> None:
+        self._input = input
+
+    def log_model(self, model: Model) -> None:
+        self._models.append(model)
+
+    def log_output(self, output: Any) -> Trace:
+        self._output = output
+
+        delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
+        self._trace = Trace(
+            created=self._start_time.isoformat(),
+            execution_time_ms=delta_time,
+            input=self._input,
+            models=self._models,
+            output=self._output,
+        )
+        return self._trace
+
+    @classmethod
+    def get_current_context(cls) -> Optional["TracingContext"]:
+        if cls._contexts[threading.get_ident()]:
+            return cls._contexts[threading.get_ident()][-1]
+        return None
+
+    def __enter__(self) -> "TracingContext":
+        self._contexts[threading.get_ident()].append(self)
+        return self
+
+    def __exit__(
+        self,
+        type: Optional[Type[BaseException]],
+        exception: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        assert self._contexts[threading.get_ident()][-1] == self
+        self._contexts[threading.get_ident()].remove(self)
+
+        if type is None:
+            assert self._trace is not None
+            self.persistence_driver.save_document(self._trace.dict())
+        else:
+            logger.exception(f"Could not finish operation: {exception}")
+
+        return True

good_ai/src/good_ai/open_s3/__main__.py

@@ -3,8 +3,8 @@
 import logging
 from pathlib import Path
 
-from large_file import LargeFile
+from .large_file import LargeFile
-from parse_arguments import parse_arguments
+from .parse_arguments import parse_arguments
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)

good_ai/src/good_ai/open_s3/helper/bytes_to_megabytes.py

@@ -0,0 +1,2 @@
+def bytes_to_megabytes(bytes: int) -> str:
+    return f"{round(bytes / 1000 / 1000, 2):.2f}"

good_ai/src/good_ai/open_s3/helper/progress_bar.py

@@ -3,6 +3,8 @@ import threading
 from logging import Logger
 from pathlib import Path
 
+from .bytes_to_megabytes import bytes_to_megabytes
+
 
 class ProgressBar:
     def __init__(self, file_size: int, logger: Logger, prefix: str):
@@ -17,10 +19,12 @@ class ProgressBar:
         with self._lock:
             self._seen_so_far += bytes_amount
             percentage = (self._seen_so_far / float(self._file_size)) * 100
-            size_length = len(str(self._file_size))
+
-            progress = str(self._seen_so_far).rjust(size_length)
+            file_size_mb = bytes_to_megabytes(self._file_size)
+            seen_so_far_mb = bytes_to_megabytes(self._seen_so_far)
+            progress = seen_so_far_mb.rjust(len(file_size_mb))
             self._logger.info(
-                f"{self._prefix} {progress}/{self._file_size} bytes ({percentage:.1f}%)"
+                f"{self._prefix} {progress}/{file_size_mb} MB ({percentage:.1f}%)"
             )
 
 

good_ai/src/good_ai/open_s3/large_file.py

@@ -5,7 +5,7 @@ import shutil
 import tempfile
 from pathlib import Path
 from types import TracebackType
-from typing import IO, Any, Dict, List, Optional, Type, Union
+from typing import IO, Any, List, Mapping, Optional, Type, Union, cast
 
 import boto3
 
@@ -65,7 +65,7 @@ class LargeFile:
         offline_mode: bool = False,
     ):
         self._name: str = name
-        self._version = version
+        self._version: int = cast(int, version)
         self._mode: str = mode
         self._keep_last_n = keep_last_n
         self._offline_mode = offline_mode
@@ -89,7 +89,7 @@ class LargeFile:
         aws_secret_access_key: str,
         large_files_bucket_name: str,
         endpoint_url: Optional[str] = None,
-        **_: Dict[str, Any],
+        **_: Mapping[str, Any],
     ) -> None:
         cls.region_name = aws_region_name
         cls.access_key_id = aws_access_key_id
@@ -188,7 +188,7 @@ class LargeFile:
                     Filename=str(tmp_file_archive),
                     Callback=None
                     if hide_progress
-                    else DownloadProgressBar(size=size, name=key, logger=logger),
+                    else DownloadProgressBar(name=str(key), size=size, logger=logger),
                 )
                 logger.info(f"Decompressing {self._local_name}")
                 shutil.unpack_archive(str(tmp_file_archive), tmp, "gztar")
@@ -199,7 +199,7 @@ class LargeFile:
 
         return destination
 
-    def push(self, path: Union[Path, str], hide_progress: bool = False) -> int:
+    def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
         if isinstance(path, str):
             path = Path(path)
 
@@ -235,7 +235,7 @@ class LargeFile:
                 Key=self._s3_name,
                 Callback=None
                 if hide_progress
-                else UploadProgressBar(file_to_be_uploaded, logger=logger),
+                else UploadProgressBar(path=file_to_be_uploaded, logger=logger),
             )
 
         self.clean_up()

good_ai/tests/open_s3/test_human_readable_to_byte.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.open_s3.helper import human_readable_to_byte
+from src.good_ai.open_s3.helper import human_readable_to_byte
 
 
 class TestHumanReadableToByte(unittest.TestCase):

good_ai/tests/open_s3/test_large_file.py

@@ -8,7 +8,7 @@ import botocore.session
 PATH = Path(__file__).parent.resolve()
 
 
-from src.open_s3 import LargeFile
+from src.good_ai import LargeFile
 
 credentials = {
     "aws_region_name": "your_region_like_eu-west-2",
@@ -30,7 +30,7 @@ class TestLargeFile(unittest.TestCase):
         self.assertRaises(ValueError, LargeFile, "test-file", "test")
 
     @patch("botocore.session")
-    def test_initialized_with_dict(self, session) -> None:
+    def test_initialized_with_dict(self, session: Any) -> None:
         session_mock = Mock()
         session.get_session = create_autospec(
             botocore.session.get_session, return_value=session_mock
@@ -58,7 +58,13 @@ class TestLargeFile(unittest.TestCase):
 
         session_mock.create_client = Mock(return_value=s3)
 
-        LargeFile.configure_credentials(**credentials)
+        LargeFile.configure_credentials(
+            aws_region_name=credentials["aws_region_name"],
+            aws_access_key_id=credentials["aws_access_key_id"],
+            aws_secret_access_key=credentials["aws_secret_access_key"],
+            large_files_bucket_name=credentials["large_files_bucket_name"],
+            endpoint_url=credentials["endpoint_url"],
+        )
         lf = LargeFile("test-file")
         session_mock.set_credentials.assert_called_once_with(
             access_key=credentials["aws_access_key_id"],

good_ai/tests/sus/data/parsed.py

@@ -1,4 +1,4 @@
-from sus.publication_tei.models import (
+from good_ai.utilities.publication_tei.models import (
     Affiliation,
     Author,
     Meta,

good_ai/tests/sus/test_clean.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.sus.clean import clean
+from src.good_ai.utilities.clean import clean
 
 
 class TestClean(unittest.TestCase):

good_ai/tests/sus/test_evaluate_ranking.py

@@ -1,7 +1,7 @@
 import unittest
 from pathlib import Path
 
-from src.sus.evaluate_ranking import evaluate_ranking
+from src.good_ai.utilities.evaluate_ranking import evaluate_ranking
 
 
 class TestEvaluateRanking(unittest.TestCase):

good_ai/tests/sus/test_get_sentences.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.sus.get_sentences import get_sentences
+from src.good_ai.utilities.get_sentences import get_sentences
 
 
 class TestGetSentences(unittest.TestCase):

good_ai/tests/sus/test_language.py

@@ -1,6 +1,10 @@
 import unittest
 
-from src.sus.language import english_name_of_language, is_english, predict_language
+from src.good_ai.utilities.language import (
+    english_name_of_language,
+    is_english,
+    predict_language,
+)
 
 
 class TestLanguage(unittest.TestCase):

good_ai/tests/sus/test_lemmatize_text.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.sus.lemmatize_text import lemmatize_text
+from src.good_ai.utilities.lemmatize_text import lemmatize_text
 
 
 class TestLemmatizeText(unittest.TestCase):

good_ai/tests/sus/test_lemmatize_token.py

@@ -1,7 +1,7 @@
 import unittest
 
-from src.sus.lemmatize_text import lemmatize_token
+from src.good_ai.utilities.lemmatize_text import lemmatize_token
-from src.sus.nlp import nlp
+from src.good_ai.utilities.nlp import nlp
 
 
 class TestLemmatizeToken(unittest.TestCase):

good_ai/tests/sus/test_match_names.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.sus.match_names.match_names import match_names
+from src.good_ai.utilities.match_names.match_names import match_names
 
 
 class TestMatchNames(unittest.TestCase):

good_ai/tests/sus/test_parallel_map.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.sus.parallel_map import parallel_map
+from src.good_ai.utilities.parallel_map import parallel_map
 
 COUNT = int(1e5) + 3
 

good_ai/tests/sus/test_publication_tei.py

@@ -1,7 +1,7 @@
 import unittest
 from pathlib import Path
 
-from src.sus.publication_tei import PublicationTEI
+from src.good_ai.utilities.publication_tei import PublicationTEI
 
 from .data.parsed import authors, content, metadata, sentences
 

good_ai/tests/sus/test_unique.py

@@ -1,6 +1,6 @@
 import unittest
 
-from src.sus.unique import unique
+from src.good_ai.utilities.unique import unique
 
 original = [
     ("a", 1),