mirror of
https://github.com/1Panel-dev/MaxKB.git
synced 2025-12-30 17:52:48 +00:00
103 lines
4.5 KiB
Python
103 lines
4.5 KiB
Python
# coding=utf-8
|
||
"""
|
||
@project: MaxKB
|
||
@Author：虎
|
||
@file： reranker.py.py
|
||
@date：2024/9/2 16:42
|
||
@desc:
|
||
"""
|
||
from typing import Sequence, Optional, Dict, Any, ClassVar
|
||
|
||
import requests
|
||
import torch
|
||
from langchain_core.callbacks import Callbacks
|
||
from langchain_core.documents import BaseDocumentCompressor, Document
|
||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||
|
||
from models_provider.base_model_provider import MaxKBBaseModel
|
||
from maxkb.const import CONFIG
|
||
|
||
|
||
class LocalReranker(MaxKBBaseModel):
|
||
def __init__(self, model_name, top_n=3, cache_dir=None):
|
||
super().__init__()
|
||
self.model_name = model_name
|
||
self.cache_dir = cache_dir
|
||
self.top_n = top_n
|
||
|
||
@staticmethod
|
||
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
|
||
if model_kwargs.get('use_local', True):
|
||
return LocalBaseReranker(model_name=model_name, cache_dir=model_credential.get('cache_dir'),
|
||
model_kwargs={'device': model_credential.get('device', 'cpu')}
|
||
|
||
)
|
||
return WebLocalBaseReranker(model_name=model_name, cache_dir=model_credential.get('cache_dir'),
|
||
model_kwargs={'device': model_credential.get('device')},
|
||
**model_kwargs)
|
||
|
||
|
||
class WebLocalBaseReranker(MaxKBBaseModel, BaseDocumentCompressor):
|
||
@staticmethod
|
||
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
|
||
pass
|
||
|
||
model_id: str = None
|
||
|
||
def __init__(self, **kwargs):
|
||
super().__init__(**kwargs)
|
||
self.model_id = kwargs.get('model_id', None)
|
||
|
||
def compress_documents(self, documents: Sequence[Document], query: str, callbacks: Optional[Callbacks] = None) -> \
|
||
Sequence[Document]:
|
||
if documents is None or len(documents) == 0:
|
||
return []
|
||
prefix = CONFIG.get_admin_path()
|
||
bind = f'{CONFIG.get("LOCAL_MODEL_HOST")}:{CONFIG.get("LOCAL_MODEL_PORT")}'
|
||
res = requests.post(
|
||
f'{CONFIG.get("LOCAL_MODEL_PROTOCOL")}://{bind}{prefix}/api/model/{self.model_id}/compress_documents',
|
||
json={'documents': [{'page_content': document.page_content, 'metadata': document.metadata} for document in
|
||
documents], 'query': query}, headers={'Content-Type': 'application/json'})
|
||
result = res.json()
|
||
if result.get('code', 500) == 200:
|
||
return [Document(page_content=document.get('page_content'), metadata=document.get('metadata')) for document
|
||
in result.get('data')]
|
||
raise Exception(result.get('message'))
|
||
|
||
|
||
class LocalBaseReranker(MaxKBBaseModel, BaseDocumentCompressor):
|
||
client: Any = None
|
||
tokenizer: Any = None
|
||
model: Optional[str] = None
|
||
cache_dir: Optional[str] = None
|
||
model_kwargs: Any = {}
|
||
|
||
def __init__(self, model_name, cache_dir=None, **model_kwargs):
|
||
super().__init__()
|
||
self.model = model_name
|
||
self.cache_dir = cache_dir
|
||
self.model_kwargs = model_kwargs
|
||
self.client = AutoModelForSequenceClassification.from_pretrained(self.model, cache_dir=self.cache_dir)
|
||
self.tokenizer = AutoTokenizer.from_pretrained(self.model, cache_dir=self.cache_dir)
|
||
self.client = self.client.to(self.model_kwargs.get('device', 'cpu'))
|
||
self.client.eval()
|
||
|
||
@staticmethod
|
||
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
|
||
return LocalBaseReranker(model_name, cache_dir=model_credential.get('cache_dir'), **model_kwargs)
|
||
|
||
def compress_documents(self, documents: Sequence[Document], query: str, callbacks: Optional[Callbacks] = None) -> \
|
||
Sequence[Document]:
|
||
if documents is None or len(documents) == 0:
|
||
return []
|
||
with torch.no_grad():
|
||
inputs = self.tokenizer([[query, document.page_content] for document in documents], padding=True,
|
||
truncation=True, return_tensors='pt', max_length=512)
|
||
scores = [torch.sigmoid(s).float().item() for s in
|
||
self.client(**inputs, return_dict=True).logits.view(-1, ).float()]
|
||
result = [Document(page_content=documents[index].page_content, metadata={'relevance_score': scores[index]})
|
||
for index
|
||
in range(len(documents))]
|
||
result.sort(key=lambda row: row.metadata.get('relevance_score'), reverse=True)
|
||
return result
|