File size: 7,128 Bytes
2999f6e 7f5a78c 2999f6e 7f5a78c 2999f6e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
import asyncio
import time
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
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 = .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
gpu_time_out = 5 # gpu processing timeout threshold
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 # 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}")
app = FastAPI(
title="baai m3, serving embed and rerank",
description="Swagger UI at https://mikeee-baai-m3.hf.space/docs",
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
@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)
|