|
import asyncio |
|
import os |
|
import time |
|
from pathlib import Path |
|
from concurrent.futures import ThreadPoolExecutor |
|
from typing import List, Tuple, Union |
|
from uuid import uuid4 |
|
|
|
from fastapi import FastAPI, HTTPException, Request |
|
from fastapi.responses import JSONResponse |
|
from FlagEmbedding import BGEM3FlagModel |
|
from pydantic import BaseModel |
|
from starlette.status import HTTP_504_GATEWAY_TIMEOUT |
|
|
|
Path("/tmp/cache").mkdir(exist_ok=True) |
|
os.environ["HF_HOME"] = "/tmp/cache" |
|
|
|
batch_size = 2 |
|
max_request = 10 |
|
max_length = 5000 |
|
max_q_length = 256 |
|
request_flush_timeout = .1 |
|
rerank_weights = [0.4, 0.2, 0.4] |
|
request_time_out = 30 |
|
gpu_time_out = 5 |
|
port= 3000 |
|
port= 7860 |
|
|
|
class m3Wrapper: |
|
def __init__(self, model_name: str, device: str = 'cuda'): |
|
"""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 |
|
self.executor = ThreadPoolExecutor() |
|
self.gpu_lock = asyncio.Semaphore(1) |
|
|
|
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: |
|
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: |
|
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}") |
|
|
|
app = FastAPI( |
|
title="baai m3, serving embed and rerank", |
|
description="Swagger UI at https://mikeee-baai-m3.hf.space/docs", |
|
version="0.1.0a0", |
|
) |
|
|
|
|
|
model = m3Wrapper('BAAI/bge-m3') |
|
processor = RequestProcessor(model, accumulation_timeout= request_flush_timeout, max_request_to_flush= max_request) |
|
|
|
|
|
@app.middleware("http") |
|
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) |
|
|
|
@app.get("/") |
|
async def landing(): |
|
"""Define landing page.""" |
|
return "Swagger UI at https://mikeee-baai-m3.hf.space/docs" |
|
|
|
@app.post("/embeddings/", response_model=EmbedResponse) |
|
async def get_embeddings(request: EmbedRequest): |
|
embeddings = await processor.process_request(request, 'embed') |
|
return EmbedResponse(embeddings=embeddings) |
|
|
|
@app.post("/rerank/", response_model=RerankResponse) |
|
async def rerank(request: RerankRequest): |
|
scores = await processor.process_request(request, 'rerank') |
|
return RerankResponse(scores=scores) |
|
|
|
if __name__ == "__main__": |
|
import uvicorn |
|
uvicorn.run(app, host="0.0.0.0", port= port) |
|
|