pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

make it lighten (#2577)

Browse Source 
### What problem does this PR solve?

#2295

### Type of change

- [x] Refactoring
This commit is contained in:
Kevin Hu
2024-09-25 13:38:40 +08:00
committed by GitHub
parent e4c9cf2264
commit dda1367ab2
4 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
rag/llm/embedding_model.py
View File

@ -245,8 +245,8 @@ class FastEmbed(Base):
threads: Optional[int] = None,
**kwargs,
):
from fastembed import TextEmbedding
if not FastEmbed._model:
if not LIGHTEN and not FastEmbed._model:
from fastembed import TextEmbedding
self._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
def encode(self, texts: list, batch_size=32):
@ -291,8 +291,8 @@ class YoudaoEmbed(Base):
_client = None
def __init__(self, key=None, model_name="maidalun1020/bce-embedding-base_v1", **kwargs):
from BCEmbedding import EmbeddingModel as qanthing
if not YoudaoEmbed._client:
if not LIGHTEN and not YoudaoEmbed._client:
from BCEmbedding import EmbeddingModel as qanthing
try:
print("LOADING BCE...")
YoudaoEmbed._client = qanthing(model_name_or_path=os.path.join(

4
rag/llm/rerank_model.py
View File

@ -109,8 +109,8 @@ class YoudaoRerank(DefaultRerank):
_model_lock = threading.Lock()
def __init__(self, key=None, model_name="maidalun1020/bce-reranker-base_v1", **kwargs):
from BCEmbedding import RerankerModel
if not YoudaoRerank._model:
if not LIGHTEN and not YoudaoRerank._model:
from BCEmbedding import RerankerModel
with YoudaoRerank._model_lock:
if not YoudaoRerank._model:
try: