Update Dockerfile m3_server.py (port = 7860)

Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.10
+ENV PIP_ROOT_USER_ACTION=ignore \
+  TZ=Asia/Shanghai
+
+WORKDIR /app
+COPY . .
+
+RUN pip install --no-cache-dir --upgrade pip && \
+  pip install --no-cache-dir -r requirements.txt
+
+# EXPOSE 7860
+# ENV PYTHONUNBUFFERED=1 \
+# 	GRADIO_ALLOW_FLAGGING=never \
+# 	GRADIO_NUM_PORTS=1 \
+# 	GRADIO_SERVER_NAME=0.0.0.0 \
+# 	GRADIO_THEME=huggingface \
+# 	SYSTEM=spaces \
+# 	SHELL=/bin/bash
+
+# CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+# CMD ["TRANSFORMERS_CACHE=./", "infinity_emb", "--model-name-or-path", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", "--port", "7860"]
+# CMD ["python", "app.py"]
+
+# CMD ["sh", "start-m3-server.sh"]
+
+CMD ["python", "m3_server.py"]

README.md

@@ -1,11 +1,12 @@
 ---
-title: Baai M3
+title: baai-m3
 emoji: 💻
 colorFrom: red
 colorTo: gray
 sdk: docker
-pinned: false
+pinned: true
 license: mit
 ---
+Swagger UI at https://mikeee-baai-m3.hf.space/docs
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

m3_server.py

@@ -0,0 +1,158 @@
+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()
+
+# 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.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)

requirements.txt

@@ -0,0 +1,3 @@
+fastapi
+flagembedding
+uvicorn

start-m3-server.sh

@@ -0,0 +1,3 @@
+export HF_HOME=/tmp/cache
+export TRANSFORMERS_CACHE=/tmp/cache
+python m3_server.py