Spaces:
Running
Running
| import asyncio | |
| import os | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from textwrap import dedent | |
| from typing import List, Tuple, Union | |
| from uuid import uuid4 | |
| import torch | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.responses import JSONResponse | |
| from FlagEmbedding import BGEM3FlagModel | |
| from pydantic import BaseModel | |
| from starlette.responses import RedirectResponse | |
| from starlette.status import HTTP_504_GATEWAY_TIMEOUT | |
| _ = """ | |
| Path("/tmp/cache").mkdir(exist_ok=True) | |
| os.environ["HF_HOME"] = "/tmp/cache" | |
| os.environ["TRANSFORMERS_CACHE"] = "/tmp/cache" | |
| # does not quite work | |
| # """ | |
| batch_size = 2 # gpu batch_size in order of your available vram | |
| max_request = 10 # max request for future improvements on api calls / gpu batches (for now is pretty basic) | |
| max_length = 5000 # max context length for embeddings and passages in re-ranker | |
| max_q_length = 256 # max context lenght for questions in re-ranker | |
| request_flush_timeout = 0.1 # flush time out for future improvements on api calls / gpu batches (for now is pretty basic) | |
| rerank_weights = [0.4, 0.2, 0.4] # re-rank score weights | |
| request_time_out = 30 # Timeout threshold | |
| request_time_out = 1200 # Timeout threshold | |
| gpu_time_out = 5 # gpu processing timeout threshold | |
| gpu_time_out = 600 # gpu processing timeout threshold | |
| port = 3000 | |
| port = 7860 | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| os.environ["TZ"] = "Asia/Shanghai" | |
| try: | |
| time.tzset() # type: ignore # pylint: disable=no-member | |
| except Exception: | |
| # Windows | |
| print("Windows, cant run time.tzset()") | |
| class m3Wrapper: | |
| def __init__(self, model_name: str, device: str = DEVICE): | |
| """Init.""" | |
| self.model = BGEM3FlagModel( | |
| model_name, device=device, use_fp16=True if device != "cpu" else False | |
| ) | |
| def embed(self, sentences: List[str]) -> List[List[float]]: | |
| embeddings = self.model.encode( | |
| sentences, batch_size=batch_size, max_length=max_length | |
| )["dense_vecs"] | |
| embeddings = embeddings.tolist() | |
| return embeddings | |
| def rerank(self, sentence_pairs: List[Tuple[str, str]]) -> List[float]: | |
| scores = self.model.compute_score( | |
| sentence_pairs, | |
| batch_size=batch_size, | |
| max_query_length=max_q_length, | |
| max_passage_length=max_length, | |
| weights_for_different_modes=rerank_weights, | |
| )["colbert+sparse+dense"] | |
| return scores | |
| class EmbedRequest(BaseModel): | |
| sentences: List[str] | |
| class RerankRequest(BaseModel): | |
| sentence_pairs: List[Tuple[str, str]] | |
| class EmbedResponse(BaseModel): | |
| embeddings: List[List[float]] | |
| class RerankResponse(BaseModel): | |
| scores: List[float] | |
| class RequestProcessor: | |
| def __init__( | |
| self, model: m3Wrapper, max_request_to_flush: int, accumulation_timeout: float | |
| ): | |
| """Init.""" | |
| self.model = model | |
| self.max_batch_size = max_request_to_flush | |
| self.accumulation_timeout = accumulation_timeout | |
| self.queue = asyncio.Queue() | |
| self.response_futures = {} | |
| self.processing_loop_task = None | |
| self.processing_loop_started = False # Processing pool flag lazy init state | |
| self.executor = ThreadPoolExecutor() # Thread pool | |
| self.gpu_lock = asyncio.Semaphore(1) # Sem for gpu sync usage | |
| async def ensure_processing_loop_started(self): | |
| if not self.processing_loop_started: | |
| print("starting processing_loop") | |
| self.processing_loop_task = asyncio.create_task(self.processing_loop()) | |
| self.processing_loop_started = True | |
| async def processing_loop(self): | |
| while True: | |
| requests, request_types, request_ids = [], [], [] | |
| start_time = asyncio.get_event_loop().time() | |
| while len(requests) < self.max_batch_size: | |
| timeout = self.accumulation_timeout - ( | |
| asyncio.get_event_loop().time() - start_time | |
| ) | |
| if timeout <= 0: | |
| break | |
| try: | |
| req_data, req_type, req_id = await asyncio.wait_for( | |
| self.queue.get(), timeout=timeout | |
| ) | |
| requests.append(req_data) | |
| request_types.append(req_type) | |
| request_ids.append(req_id) | |
| except asyncio.TimeoutError: | |
| break | |
| if requests: | |
| await self.process_requests_by_type( | |
| requests, request_types, request_ids | |
| ) | |
| async def process_requests_by_type(self, requests, request_types, request_ids): | |
| tasks = [] | |
| for request_data, request_type, request_id in zip( | |
| requests, request_types, request_ids | |
| ): | |
| if request_type == "embed": | |
| task = asyncio.create_task( | |
| self.run_with_semaphore( | |
| self.model.embed, request_data.sentences, request_id | |
| ) | |
| ) | |
| else: # 'rerank' | |
| task = asyncio.create_task( | |
| self.run_with_semaphore( | |
| self.model.rerank, request_data.sentence_pairs, request_id | |
| ) | |
| ) | |
| tasks.append(task) | |
| await asyncio.gather(*tasks) | |
| async def run_with_semaphore(self, func, data, request_id): | |
| async with self.gpu_lock: # Wait for sem | |
| future = self.executor.submit(func, data) | |
| try: | |
| result = await asyncio.wait_for( | |
| asyncio.wrap_future(future), timeout=gpu_time_out | |
| ) | |
| self.response_futures[request_id].set_result(result) | |
| except asyncio.TimeoutError: | |
| self.response_futures[request_id].set_exception( | |
| TimeoutError("GPU processing timeout") | |
| ) | |
| except Exception as e: | |
| self.response_futures[request_id].set_exception(e) | |
| async def process_request( | |
| self, request_data: Union[EmbedRequest, RerankRequest], request_type: str | |
| ): | |
| try: | |
| await self.ensure_processing_loop_started() | |
| request_id = str(uuid4()) | |
| self.response_futures[request_id] = asyncio.Future() | |
| await self.queue.put((request_data, request_type, request_id)) | |
| return await self.response_futures[request_id] | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}") | |
| description = dedent( | |
| r""" | |
| embed example: | |
| ```bash | |
| curl -X 'POST' \ | |
| 'https://mikeee-baai-m3.hf.space/embed/' \ | |
| -H 'accept: application/json' \ | |
| -H 'Content-Type: application/json' \ | |
| -d '{ | |
| "sentences": [ | |
| "string", "string1" | |
| ] | |
| }' | |
| ``` | |
| rerank example: | |
| ```bash | |
| ... | |
| ``` | |
| """ | |
| ) | |
| app = FastAPI( | |
| title="Serving BAAI/bge-m3 embed and rerank", | |
| # description="Swagger UI at https://mikeee-baai-m3.hf.space/docs", | |
| description=description, | |
| version="0.1.0a0", | |
| ) | |
| # Initialize the model and request processor | |
| model = m3Wrapper("BAAI/bge-m3") | |
| processor = RequestProcessor( | |
| model, accumulation_timeout=request_flush_timeout, max_request_to_flush=max_request | |
| ) | |
| # Adding a middleware returning a 504 error if the request processing time is above a certain threshold | |
| async def timeout_middleware(request: Request, call_next): | |
| try: | |
| start_time = time.time() | |
| return await asyncio.wait_for(call_next(request), timeout=request_time_out) | |
| except asyncio.TimeoutError: | |
| process_time = time.time() - start_time | |
| return JSONResponse( | |
| { | |
| "detail": "Request processing time excedeed limit", | |
| "processing_time": process_time, | |
| }, | |
| status_code=HTTP_504_GATEWAY_TIMEOUT, | |
| ) | |
| async def landing(): | |
| """Define landing page.""" | |
| # return "Swagger UI at https://mikeee-baai-m3.hf.space/docs" | |
| return RedirectResponse("/docs") | |
| async def get_embeddings(request: EmbedRequest): | |
| embeddings = await processor.process_request(request, "embed") | |
| return EmbedResponse(embeddings=embeddings) | |
| async def rerank(request: RerankRequest): | |
| scores = await processor.process_request(request, "rerank") | |
| return RerankResponse(scores=scores) | |
| if __name__ == "__main__": | |
| import uvicorn | |
| print("started") | |
| uvicorn.run(app, host="0.0.0.0", port=port) | |