File size: 3,403 Bytes
750b88d | 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 | #!/usr/bin/env python3
"""Loopback-only CPU BGE v2-m3 service matching MemOS HTTP reranking."""
import argparse
import os
import threading
import time
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--model-path', required=True)
parser.add_argument('--port', type=int, default=18086)
parser.add_argument('--threads', type=int, default=8)
parser.add_argument('--batch-size', type=int, default=8)
parser.add_argument('--max-length', type=int, default=1024)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from transformers import AutoModelForSequenceClassification, AutoTokenizer
torch.set_num_threads(args.threads)
torch.set_num_interop_threads(1)
tokenizer = AutoTokenizer.from_pretrained(args.model_path, local_files_only=True)
model = AutoModelForSequenceClassification.from_pretrained(
args.model_path, local_files_only=True).to('cpu').eval()
lock = threading.Lock()
stats = {'requests': 0, 'documents': 0, 'seconds': 0.0}
app = FastAPI(title='MemOS CPU BGE reranker')
class Request(BaseModel):
model: str = 'bge-reranker-v2-m3'
query: str
documents: list[str] = Field(max_length=1000)
@app.get('/health')
def health():
return {'status': 'ok', 'model': 'bge-reranker-v2-m3', 'device': 'cpu',
'score': 'sigmoid(raw_sequence_classification_logit)',
'max_length': args.max_length, 'batch_size': args.batch_size,
'threads': args.threads, **stats}
@app.post('/rerank')
def rerank(request: Request):
if request.model not in ('bge-reranker-v2-m3', 'BAAI/bge-reranker-v2-m3'):
raise HTTPException(400, 'Unsupported reranker model')
if not request.documents:
return {'results': []}
started = time.monotonic()
scores = []
with lock, torch.inference_mode():
try:
for offset in range(0, len(request.documents), args.batch_size):
pairs = [[request.query, doc] for doc in request.documents[offset:offset+args.batch_size]]
encoded = tokenizer(pairs, padding=True, truncation=True,
return_tensors='pt', max_length=args.max_length)
logits = model(**encoded, return_dict=True).logits.view(-1).float()
scores.extend(torch.sigmoid(logits).tolist())
stats['requests'] += 1
stats['documents'] += len(request.documents)
stats['seconds'] += time.monotonic() - started
except Exception as error:
raise HTTPException(500, f'Reranking failed: {type(error).__name__}') from None
if len(scores) != len(request.documents):
raise HTTPException(500, 'Reranker output length mismatch')
return {'results': sorted(
[{'index': i, 'relevance_score': score} for i, score in enumerate(scores)],
key=lambda row: (-row['relevance_score'], row['index']))}
uvicorn.run(app, host='127.0.0.1', port=args.port, access_log=False, log_level='warning')
if __name__ == '__main__':
main()
|