Просмотр исходного кода

Let parallism of RAPTOR controlable. (#5379)

### What problem does this PR solve?

#4874
### Type of change

- [x] Refactoring
tags/v0.17.0
Kevin Hu 8 месяцев назад
Родитель
Сommit
96e9d50060
Аккаунт пользователя с таким Email не найден
1 измененных файлов: 2 добавлений и 1 удалений
  1. 2
    1
      rag/raptor.py

+ 2
- 1
rag/raptor.py Просмотреть файл

@@ -14,6 +14,7 @@
# limitations under the License.
#
import logging
import os
import re
from concurrent.futures import ThreadPoolExecutor, ALL_COMPLETED, wait
from threading import Lock
@@ -122,7 +123,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
lbls = [np.where(prob > self._threshold)[0] for prob in probs]
lbls = [lbl[0] if isinstance(lbl, np.ndarray) else lbl for lbl in lbls]
lock = Lock()
with ThreadPoolExecutor(max_workers=12) as executor:
with ThreadPoolExecutor(max_workers=int(os.environ.get('GRAPH_EXTRACTOR_MAX_WORKERS', 10))) as executor:
threads = []
for c in range(n_clusters):
ck_idx = [i + start for i in range(len(lbls)) if lbls[i] == c]

Загрузка…
Отмена
Сохранить