Add auto-reload, evaluation endpoint, show docs

.vscode/settings.json

@@ -2,6 +2,7 @@
   "cSpell.words": [
     "boto",
     "botocore",
+    "fastapi",
     "iloc",
     "inplace",
     "levelno",

.vscode/tasks.json

@@ -4,9 +4,9 @@
     {
       "label": "Format and lint Python",
       "type": "shell",
-      "command": "source .env/bin/activate && scripts/format-python.sh great_ai && scripts/format-python.sh example",
+      "command": "source .env/bin/activate && scripts/format-python.sh great_ai && scripts/format-python.sh examples",
       "windows": {
-        "command": ".env/bin/activate.bat; scripts/format-python.sh great_ai; scripts\\format-python.sh example"
+        "command": ".env/bin/activate.bat; scripts/format-python.sh great_ai; scripts\\format-python.sh examples"
       },
       "group": "test",
       "presentation": {

example/main_service.py

@@ -1,9 +0,0 @@
-#!/usr/bin/env python3
-
-
-from predict_domain import predict_domain
-
-from great_ai import serve
-
-if __name__ == "__main__":
-    serve(predict_domain)

examples/README.md

@@ -0,0 +1,12 @@
+# Train Domain classifier on 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 great_ai.open_s3 --secrets s3.ini --push ss-data
+rm -rf ss-data
+```

examples/main_batch.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+
+import json
+from random import shuffle
+
+from devtools import debug
+from predict_domain import predict_domain
+
+from great_ai import process_batch
+
+if __name__ == "__main__":
+    with open(".cache/data-1/s2-corpus-323.json") as f:
+        raw = json.load(f)
+
+    shuffle(raw)
+    data = {f'{r["title"]} {r["abstract"]}': r["domain"] for r in raw[:10]}
+
+    results = process_batch(predict_domain, data.keys())
+
+    for predicted, actual in zip(results, data.values()):
+        print(", ".join(actual))
+        debug(predicted)
+        print()

examples/main_service.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python3
+
+from great_ai import configure, create_service
+
+configure(development_mode_override=True)
+
+from predict_domain import predict_domain
+
+app = create_service(predict_domain)

examples/predict_domain.py

@@ -0,0 +1,79 @@
+import re
+from typing import Dict, Iterable, List
+
+from preprocess import preprocess
+from pydantic import BaseModel
+from sklearn.pipeline import Pipeline
+
+from great_ai import log_argument, log_metric, use_model
+from great_ai.utilities.clean import clean
+
+
+class DomainPrediction(BaseModel):
+    domain: str
+    probability: float
+    explanation: List[str]
+
+
+@use_model("small-domain-prediction-v2", version="latest")
+@log_argument("text", validator=lambda t: len(t) > 0)
+def predict_domain(
+    text: str, model: Pipeline, cut_off_probability: float = 0.2
+) -> List[DomainPrediction]:
+    """
+    Predict the scientific domain of the input text.
+    Return labels until their sum likelihood is larger than cut_off_probability.
+    """
+    log_metric("text_length", len(text))
+
+    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 * 100:
+            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]

examples/preprocess.py

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

great_ai/src/great_ai/__main__.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+import re
+from importlib import import_module
+from sys import argv
+
+import uvicorn
+from uvicorn.config import LOGGING_CONFIG
+
+from great_ai.great_ai.context import get_context
+
+if __name__ == "__main__":
+    if len(argv) < 2:
+        raise ValueError(f"Provide a filename such as: {argv[0]} my_app.py")
+
+    module_name = re.sub(r"\.py\b", "", argv[1])
+    if ":" not in module_name:
+        module_name += ":app"
+
+    base = module_name.split(":")[0]
+    module = import_module(base)
+    get_context().logger.info("Starting server")
+
+    uvicorn.run(
+        module_name,
+        host="0.0.0.0",
+        port=6060,
+        workers=1,
+        reload=not get_context().is_production,
+        log_config={
+            **LOGGING_CONFIG,
+            "formatters": {
+                "default": {
+                    "()": "great_ai.logger.CustomFormatter",
+                    "fmt": "%(asctime)s | %(levelname)8s | %(message)s",
+                },
+                "access": {
+                    "()": "great_ai.logger.CustomFormatter",
+                    "fmt": "%(asctime)s | %(levelname)8s | %(message)s",  # noqa: E501
+                },
+            },
+        },
+    )

great_ai/src/great_ai/great_ai/__init__.py

@@ -1,4 +1,4 @@
-from .context import set_default_config
+from .context import configure
-from .deploy import process_batch, process_single, serve
+from .deploy import create_service, process_batch, process_single
 from .metrics import log_argument, log_metric
 from .models import save_model, use_model

great_ai/src/great_ai/great_ai/context/__init__.py

@@ -1,2 +1,2 @@
 from .context import Context
-from .get_context import get_context, set_default_config
+from .get_context import configure, get_context

great_ai/src/great_ai/great_ai/context/get_context.py

@@ -16,26 +16,35 @@ PRODUCTION_KEY = "production"
 
 def get_context() -> Context:
     if _context is None:
-        set_default_config()
+        configure()
 
     return cast(Context, _context)
 
 
-def set_default_config(
+def configure(
     log_level: int = INFO,
     s3_config: Path = Path("s3.ini"),
     seed: int = 42,
     persistence_driver: PersistenceDriver = ParallelTinyDbDriver(
         Path("tracing_database.json")
     ),
-    is_production_mode_override: Optional[bool] = None,
+    development_mode_override: Optional[bool] = None,
 ) -> None:
     global _context
 
     logger = create_logger("great_ai", level=log_level)
 
+    if _context is not None:
+        logger.warn(
+            "Configuration has been already initialised, overwriting.\n"
+            + 'Make sure to call "configure()" before importing your application code.'
+        )
+
     is_production = _is_in_production_mode(
-        override=is_production_mode_override, logger=logger
+        override=None
+        if development_mode_override is None
+        else not development_mode_override,
+        logger=logger,
     )
     _initialize_large_file(s3_config, logger=logger)
     _set_seed(seed)
@@ -51,7 +60,7 @@ def set_default_config(
         logger=logger,
     )
 
-    logger.info("Defaults: configured ✅")
+    logger.info("Options: configured ✅")
 
 
 def _is_in_production_mode(override: Optional[bool], logger: Logger) -> bool:

great_ai/src/great_ai/great_ai/deploy/__init__.py

@@ -1,3 +1,3 @@
+from .create_service import create_service
 from .process_batch import process_batch
 from .process_single import process_single
-from .serve import serve

great_ai/src/great_ai/great_ai/deploy/create_service.py

@@ -0,0 +1,94 @@
+from pathlib import Path
+from typing import Any, Callable, Dict, List
+
+from fastapi import FastAPI, HTTPException, status
+from fastapi.middleware.wsgi import WSGIMiddleware
+from fastapi.openapi.docs import get_swagger_ui_html
+from fastapi.responses import RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from starlette.responses import HTMLResponse
+
+from ..context import get_context
+from ..helper import snake_case_to_text
+from ..metrics import create_dash_app
+from ..tracing import TracingContext
+from ..views import EvaluationFeedbackRequest, HealthCheckResponse, Query, Trace
+
+PATH = Path(__file__).parent.resolve()
+
+
+def create_service(
+    func: Callable[..., Any], disable_docs: bool = False, disable_metrics: bool = False
+) -> FastAPI:
+    function_name = func.__name__
+    function_docs = func.__doc__
+
+    documentation = (
+        f"REST API wrapper for interacting with the '{function_name}' function.\n"
+    )
+    if function_docs:
+        documentation += function_docs
+
+    app = FastAPI(
+        title=snake_case_to_text(function_name),
+        description=documentation,
+        docs_url=None,
+        redoc_url=None,
+    )
+
+    @app.post("/evaluations", status_code=status.HTTP_200_OK, response_model=Trace)
+    def score(input: Any) -> Trace:
+        with TracingContext() as t:
+            result = func(input)
+            output = t.log_output(result)
+        return output
+
+    @app.get("/evaluations/:evaluation_id", status_code=status.HTTP_200_OK)
+    def get_evaluation(evaluation_id: str) -> Trace:
+        result = get_context().persistence.get_trace(evaluation_id)
+        if result is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+        return result
+
+    @app.post(
+        "/evaluations/:evaluation_id/feedback", status_code=status.HTTP_202_ACCEPTED
+    )
+    def give_feedback(evaluation_id: str, input: EvaluationFeedbackRequest) -> None:
+        get_context().persistence.add_evaluation(evaluation_id, input.evaluation)
+
+    if not disable_docs:
+
+        @app.get("/docs", include_in_schema=False)
+        def custom_swagger_ui_html() -> HTMLResponse:
+            return get_swagger_ui_html(openapi_url="openapi.json", title=app.title)
+
+        @app.get("/docs/index.html", include_in_schema=False)
+        def redirect_to_entrypoint() -> RedirectResponse:
+            return RedirectResponse("/docs")
+
+    if not disable_metrics:
+        dash_app = create_dash_app(function_name, documentation)
+        app.mount(get_context().metrics_path, WSGIMiddleware(dash_app))
+
+        @app.get("/", include_in_schema=False)
+        def redirect_to_entrypoint() -> RedirectResponse:
+            return RedirectResponse("/metrics")
+
+        app.mount(
+            "/assets", StaticFiles(directory=PATH / "../metrics/assets"), name="static"
+        )
+
+        @app.post("/query", status_code=status.HTTP_200_OK)
+        def query_metrics(query: Query) -> List[Dict[str, Any]]:
+            return get_context().persistence.query(
+                conjunctive_filters=query.filter,
+                sort_by=query.sort,
+                skip=query.skip,
+                take=query.take,
+            )
+
+    @app.get("/health", status_code=status.HTTP_200_OK)
+    def check_health() -> HealthCheckResponse:
+        return HealthCheckResponse(is_healthy=True)
+
+    return app

great_ai/src/great_ai/great_ai/deploy/serve.py

@@ -1,48 +0,0 @@
-from typing import Any, Callable
-
-import uvicorn
-from fastapi import FastAPI, status
-from uvicorn.config import LOGGING_CONFIG
-
-from ..tracing import TracingContext
-from ..views import Trace
-from .create_fastapi_app import create_fastapi_app
-
-
-def serve(
-    function: Callable[..., Any],
-    disable_docs: bool = False,
-    disable_metrics: bool = False,
-    configure: Callable[[FastAPI], None] = lambda _: None,
-) -> None:
-    app = create_fastapi_app(
-        function.__name__, disable_docs=disable_docs, disable_metrics=disable_metrics
-    )
-
-    @app.post("/score", status_code=status.HTTP_200_OK, response_model=Trace)
-    def process(input: Any) -> Trace:
-        with TracingContext() as t:
-            result = function(input)
-            output = t.log_output(result)
-        return output
-
-    configure(app)
-
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=5050,
-        log_config={
-            **LOGGING_CONFIG,
-            "formatters": {
-                "default": {
-                    "()": "great_ai.logger.CustomFormatter",
-                    "fmt": "%(asctime)s | %(levelname)8s | %(message)s",
-                },
-                "access": {
-                    "()": "great_ai.logger.CustomFormatter",
-                    "fmt": "%(asctime)s | %(levelname)8s | %(message)s",  # noqa: E501
-                },
-            },
-        },
-    )

great_ai/src/great_ai/great_ai/exceptions/__init__.py

@@ -1 +1,2 @@
 from .argument_validation_error import ArgumentValidationError
+from .missing_argument_error import MissingArgumentError

great_ai/src/great_ai/great_ai/exceptions/missing_argument_error.py

@@ -0,0 +1,2 @@
+class MissingArgumentError(Exception):
+    pass

great_ai/src/great_ai/great_ai/helper/__init__.py

@@ -1,3 +1,4 @@
 from .get_args import get_args
 from .snake_case_to_text import snake_case_to_text
+from .strip_lines import strip_lines
 from .text_to_hex_color import text_to_hex_color

great_ai/src/great_ai/great_ai/helper/strip_lines.py

@@ -0,0 +1,2 @@
+def strip_lines(text: str) -> str:
+    return "\n".join(line.strip() for line in text.split("\n"))