add lighten control (#2567)

### What problem does this PR solve?

#2295

### Type of change

- [ ] Bug Fix (non-breaking change which fixes an issue)
- [x] New Feature (non-breaking change which adds functionality)
- [ ] Documentation Update
- [ ] Refactoring
- [ ] Performance Improvement
- [ ] Other (please describe):
This commit is contained in:
Kevin Hu
2024-09-24 19:22:01 +08:00
committed by GitHub
parent 9251fb39af
commit 7bb28ca2bd
5 changed files with 86 additions and 78 deletions

View File

@ -14,14 +14,14 @@
# limitations under the License.
#
import re
import threading
import threading
import requests
import torch
from FlagEmbedding import FlagReranker
from huggingface_hub import snapshot_download
import os
from abc import ABC
import numpy as np
from api.settings import LIGHTEN
from api.utils.file_utils import get_home_cache_dir
from rag.utils import num_tokens_from_string, truncate
import json
@ -53,7 +53,9 @@ class DefaultRerank(Base):
^_-
"""
if not DefaultRerank._model:
if not LIGHTEN and not DefaultRerank._model:
import torch
from FlagEmbedding import FlagReranker
with DefaultRerank._model_lock:
if not DefaultRerank._model:
try: