23 lines
705 B
Python
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 |