mirror of
https://github.com/1Panel-dev/MaxKB.git
synced 2025-12-26 01:33:05 +00:00
fix: Rearrange model dependency (#4324)
This commit is contained in:
parent
3ab17c61ff
commit
1b464c4d8f
|
|
@ -11,7 +11,6 @@ from typing import Sequence, Optional, Dict, Any
|
|||
|
||||
from langchain_core.callbacks import Callbacks
|
||||
from langchain_core.documents import Document, BaseDocumentCompressor
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
from models_provider.base_model_provider import MaxKBBaseModel
|
||||
|
||||
|
|
@ -25,6 +24,7 @@ class LocalReranker(MaxKBBaseModel, BaseDocumentCompressor):
|
|||
|
||||
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
|
||||
|
|
|
|||
Loading…
Reference in New Issue