OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-09 16:07:54 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

refac: reranker

Browse Source 
Co-Authored-By: Tornike Gurgenidze <togurg14@freeuni.edu.ge>
This commit is contained in:
Timothy Jaeryang Baek 2025-05-23 01:29:48 +04:00
parent f2d2cae8fb
commit aac25eac9e
3 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
backend/open_webui/retrieval/models/base_reranker.py Normal file
View 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

4
backend/open_webui/retrieval/models/colbert.py
View File

@ -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"

4
backend/open_webui/retrieval/models/external.py
View File

@ -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,