mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-09-09 21:08:59 +00:00
feat(model): Support deploy rerank model (#1522)
This commit is contained in:
@@ -1,9 +1,24 @@
|
||||
"""Interface for embedding models."""
|
||||
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List
|
||||
|
||||
|
||||
class RerankEmbeddings(ABC):
|
||||
"""Interface for rerank models."""
|
||||
|
||||
@abstractmethod
|
||||
def predict(self, query: str, candidates: List[str]) -> List[float]:
|
||||
"""Predict the scores of the candidates."""
|
||||
|
||||
async def apredict(self, query: str, candidates: List[str]) -> List[float]:
|
||||
"""Asynchronously predict the scores of the candidates."""
|
||||
return await asyncio.get_running_loop().run_in_executor(
|
||||
None, self.predict, query, candidates
|
||||
)
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""Interface for embedding models.
|
||||
|
||||
|
Reference in New Issue
Block a user