diff --git a/deepsearch/cli.py b/deepsearch/cli.py index 9b7bc970..ef52bde3 100644 --- a/deepsearch/cli.py +++ b/deepsearch/cli.py @@ -4,7 +4,6 @@ from deepsearch.core.cli.plugins import get_cli_groups from deepsearch.cps.cli.main import app as cps_app from deepsearch.documents.cli.main import app as documents_app -from deepsearch.model.cli.main import app as model_app from deepsearch.query.cli.main import app as query_app app.add_typer(cps_app, name="cps", help="Interact with DeepSearch CPS component") @@ -14,9 +13,6 @@ name="documents", help="Interact with DeepSearch Document Conversion component", ) -app.add_typer( - model_app, name="model", help="Interact with the DeepSearch model component" -) for group in get_cli_groups(): app.add_typer(group) diff --git a/deepsearch/model/README.md b/deepsearch/model/README.md new file mode 100644 index 00000000..5529d302 --- /dev/null +++ b/deepsearch/model/README.md @@ -0,0 +1,24 @@ +# Model API + +> Currently in **beta**. + +The Model API allows users to serve and integrate their own models. + +## Usage +```python +from deepsearch.model.server.deepsearch_annotator_app import DeepSearchAnnotatorApp + +annotator = ... # e.g. SimpleTextGeographyAnnotator() +app = DeepSearchAnnotatorApp() +app.register_annotator(annotator) +app.run() +``` + +For a complete example, check [examples/main.py](examples/main.py). + +## Models +For an example, check +[examples/simple_text_geography_annotator/](examples/simple_text_geography_annotator/). + +## Base models +Check [base/](base/). diff --git a/deepsearch/model/cli/__init__.py b/deepsearch/model/cli/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/deepsearch/model/cli/main.py b/deepsearch/model/cli/main.py deleted file mode 100644 index c44eb616..00000000 --- a/deepsearch/model/cli/main.py +++ /dev/null @@ -1,31 +0,0 @@ -import importlib.util -import os - -import typer -import uvicorn - -from deepsearch.model.examples.simple_text_geography_annotator.simple_text_geography_annotator import ( # type: ignore - SimpleTextGeographyAnnotator, -) -from deepsearch.model.server.deepsearch_annotator_app import DeepSearchAnnotatorApp - -app = typer.Typer(no_args_is_help=True) - - -@app.command(name="serve", help="Serve an annotator instance") -def serve( - port: int = typer.Option(8000, "-p", "--port", help="The port to listen on"), - annotator: str = typer.Option(None, "-a", "--annotator", help="Not implemented"), -): - # Load the Annotator App - deepsearch_annotator_app = DeepSearchAnnotatorApp() - - # register annotators - deepsearch_annotator_app.register_annotator(SimpleTextGeographyAnnotator()) - - # Run - uvicorn.run(deepsearch_annotator_app.app, host="0.0.0.0", port=port) - - -if __name__ == "__main__": - app() diff --git a/deepsearch/model/examples/main.py b/deepsearch/model/examples/main.py new file mode 100644 index 00000000..785d8dfc --- /dev/null +++ b/deepsearch/model/examples/main.py @@ -0,0 +1,15 @@ +from deepsearch.model.examples.simple_text_geography_annotator.simple_text_geography_annotator import ( # type: ignore + SimpleTextGeographyAnnotator, +) +from deepsearch.model.server.deepsearch_annotator_app import DeepSearchAnnotatorApp + + +def run(): + app = DeepSearchAnnotatorApp() + app.register_annotator(SimpleTextGeographyAnnotator()) + app.register_annotator(SimpleTextGeographyAnnotator(), name="foo") + app.run() + + +if __name__ == "__main__": + run() diff --git a/deepsearch/model/server/deepsearch_annotator_app.py b/deepsearch/model/server/deepsearch_annotator_app.py index 2477e097..b6574cc7 100644 --- a/deepsearch/model/server/deepsearch_annotator_app.py +++ b/deepsearch/model/server/deepsearch_annotator_app.py @@ -4,12 +4,11 @@ import inspect import logging import os -import sys import time -from concurrent.futures import ThreadPoolExecutor from datetime import datetime -from typing import Coroutine, Dict, List, Optional, Union +from typing import Coroutine, Dict, Union +import uvicorn from anyio import CapacityLimiter from anyio.lowlevel import RunVar from fastapi import FastAPI, HTTPException, Request @@ -66,9 +65,12 @@ async def validation_exception_handler( async def health_check() -> dict: return {"message": "HealthCheck"} - @self.app.get("/annotator") - async def health_check() -> dict: - return {key: f"annotator/{key}" for key in self.annotators_list} + @self.app.get("/") + async def get_definitions() -> dict: + return { + key: self.annotate_controller.get_annotator_info(key) + for key in self.annotators_list + } # Will Require an API key @self.app.get("/annotator/{annotator_name}") @@ -197,6 +199,9 @@ def register_annotator( ) exit(-1) + def run(self, host: str = "127.0.0.1", port: int = 8000, **kwargs) -> None: + uvicorn.run(self.app, host=host, port=port, **kwargs) + class AnnotateController: def __init__(self, app_annotators): self.app_annotators: dict = app_annotators