| |
| """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() |
|
|
