import logging import os from pathlib import Path from typing import List, Union import psutil import torch from relik.common.utils import is_package_available from relik.inference.annotator import Relik if not is_package_available("fastapi"): raise ImportError( "FastAPI is not installed. Please install FastAPI with `pip install relik[serve]`." ) from fastapi import FastAPI, HTTPException, APIRouter from relik.common.log import get_logger from relik.inference.serve.backend.utils import ( RayParameterManager, ServerParameterManager, ) logger = get_logger(__name__, level=logging.INFO) VERSION = {} # type: ignore with open( Path(__file__).parent.parent.parent.parent / "version.py", "r" ) as version_file: exec(version_file.read(), VERSION) # Env variables for server SERVER_MANAGER = ServerParameterManager() RAY_MANAGER = RayParameterManager() class RelikServer: def __init__( self, relik_pretrained: str | None = None, device: str = "cpu", retriever_device: str | None = None, document_index_device: str | None = None, reader_device: str | None = None, precision: str | int | torch.dtype = 32, retriever_precision: str | int | torch.dtype | None = None, document_index_precision: str | int | torch.dtype | None = None, reader_precision: str | int | torch.dtype | None = None, annotation_type: str = "char", **kwargs, ): num_threads = os.getenv("TORCH_NUM_THREADS", psutil.cpu_count(logical=False)) torch.set_num_threads(num_threads) logger.info(f"Torch is running on {num_threads} threads.") # parameters logger.info(f"RELIK_PRETRAINED: {relik_pretrained}") self.relik_pretrained = relik_pretrained logger.info(f"DEVICE: {device}") self.device = device if retriever_device is not None: logger.info(f"RETRIEVER_DEVICE: {retriever_device}") self.retriever_device = retriever_device or device if document_index_device is not None: logger.info(f"INDEX_DEVICE: {document_index_device}") self.document_index_device = document_index_device or retriever_device if reader_device is not None: logger.info(f"READER_DEVICE: {reader_device}") self.reader_device = reader_device logger.info(f"PRECISION: {precision}") self.precision = precision if retriever_precision is not None: logger.info(f"RETRIEVER_PRECISION: {retriever_precision}") self.retriever_precision = retriever_precision or precision if document_index_precision is not None: logger.info(f"INDEX_PRECISION: {document_index_precision}") self.document_index_precision = document_index_precision or precision if reader_precision is not None: logger.info(f"READER_PRECISION: {reader_precision}") self.reader_precision = reader_precision or precision logger.info(f"ANNOTATION_TYPE: {annotation_type}") self.annotation_type = annotation_type self.relik = Relik.from_pretrained( self.relik_pretrained, device=self.device, retriever_device=self.retriever_device, document_index_device=self.document_index_device, reader_device=self.reader_device, precision=self.precision, retriever_precision=self.retriever_precision, document_index_precision=self.document_index_precision, reader_precision=self.reader_precision, ) self.router = APIRouter() self.router.add_api_route("/api/relik", self.relik_endpoint, methods=["POST"]) logger.info("RelikServer initialized.") # @serve.batch() async def __call__(self, text: List[str]) -> List: return self.relik(text, annotation_type=self.annotation_type) # @app.post("/api/relik") async def relik_endpoint(self, text: Union[str, List[str]]): try: # get predictions for the retriever return await self(text) except Exception as e: # log the entire stack trace logger.exception(e) raise HTTPException(status_code=500, detail=f"Server Error: {e}") app = FastAPI( title="ReLiK", version=VERSION["VERSION"], description="ReLiK REST API", ) server = RelikServer(**vars(SERVER_MANAGER)) app.include_router(server.router)