mirror of
https://github.com/open-webui/open-webui
synced 2025-06-11 17:02:47 +00:00
refac: reranker
Co-Authored-By: Tornike Gurgenidze <togurg14@freeuni.edu.ge>
This commit is contained in:
parent
f2d2cae8fb
commit
aac25eac9e
8
backend/open_webui/retrieval/models/base_reranker.py
Normal file
8
backend/open_webui/retrieval/models/base_reranker.py
Normal file
@ -0,0 +1,8 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, List, Tuple
|
||||
|
||||
|
||||
class BaseReranker(ABC):
|
||||
@abstractmethod
|
||||
def predict(self, sentences: List[Tuple[str, str]]) -> Optional[List[float]]:
|
||||
pass
|
@ -7,11 +7,13 @@ from colbert.modeling.checkpoint import Checkpoint
|
||||
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
|
||||
from open_webui.retrieval.models.base_reranker import BaseReranker
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
|
||||
class ColBERT:
|
||||
class ColBERT(BaseReranker):
|
||||
def __init__(self, name, **kwargs) -> None:
|
||||
log.info("ColBERT: Loading model", name)
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
@ -3,12 +3,14 @@ import requests
|
||||
from typing import Optional, List, Tuple
|
||||
|
||||
from open_webui.env import SRC_LOG_LEVELS
|
||||
from open_webui.retrieval.models.base_reranker import BaseReranker
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
||||
|
||||
|
||||
class ExternalReranker:
|
||||
class ExternalReranker(BaseReranker):
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
|
Loading…
Reference in New Issue
Block a user