| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758 |
- # coding=utf-8
- """
- @project: MaxKB
- @Author：虎虎
- @file： model.py
- @date：2025/11/5 15:30
- @desc:
- """
- from typing import Sequence, Optional, Dict, Any
- from langchain_core.callbacks import Callbacks
- from langchain_core.documents import Document, BaseDocumentCompressor
- from models_provider.base_model_provider import MaxKBBaseModel
- class LocalReranker(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__()
- from transformers import AutoModelForSequenceClassification, AutoTokenizer
- 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 is_cache_model():
- return False
- @staticmethod
- def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
- return LocalReranker(model_name, cache_dir=model_credential.get('cache_dir'))
- def compress_documents(self, documents: Sequence[Document], query: str, callbacks: Optional[Callbacks] = None) -> \
- Sequence[Document]:
- if documents is None or len(documents) == 0:
- return []
- import torch
- 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
|
