llmengine/llmengine/base_bgererank.py
2025-07-08 18:09:55 +08:00

23 lines
705 B
Python

from abc import ABC, abstractmethod
from typing import List, Dict
model_pathMap = {}
def rerank_register(model_key, Klass):
"""Register a rerank model class for a given model key."""
global model_pathMap
model_pathMap[model_key] = Klass
def get_rerank_class(model_path):
"""Find the rerank model class for a given model path."""
for k, klass in model_pathMap.items():
if len(model_path.split(k)) > 1:
return klass
print(f'{model_pathMap=}')
return None
class BaseRerank(ABC):
@abstractmethod
async def rerank(self, query: str, documents: List[str], top_k: int) -> List[Dict]:
"""Rerank documents based on query relevance."""
pass