diff --git a/deepdoc/vision/layout_recognizer.py b/deepdoc/vision/layout_recognizer.py index 39e6e92b5..e7ad76105 100644 --- a/deepdoc/vision/layout_recognizer.py +++ b/deepdoc/vision/layout_recognizer.py @@ -15,6 +15,8 @@ import re from collections import Counter from copy import deepcopy import numpy as np + +from api.db import ParserType from api.utils.file_utils import get_project_base_directory from deepdoc.vision import Recognizer @@ -35,6 +37,7 @@ class LayoutRecognizer(Recognizer): ] def __init__(self, domain): super().__init__(self.labels, domain, os.path.join(get_project_base_directory(), "rag/res/deepdoc/")) + self.garbage_layouts = ["footer", "header", "reference"] def __call__(self, image_list, ocr_res, scale_factor=3, thr=0.2, batch_size=16): def __is_garbage(b): @@ -85,7 +88,7 @@ class LayoutRecognizer(Recognizer): i += 1 continue lts_[ii]["visited"] = True - if lts_[ii]["type"] in ["footer", "header", "reference"]: + if lts_[ii]["type"] in self.garbage_layouts: if lts_[ii]["type"] not in garbages: garbages[lts_[ii]["type"]] = [] garbages[lts_[ii]["type"]].append(bxs[i]["text"]) diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index f1e7c1ec6..a37423506 100644 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -6,11 +6,10 @@ export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/ PY=/root/miniconda3/envs/py11/bin/python - - function task_exe(){ - sleep 60; - while [ 1 -eq 1 ];do mpirun -n 4 --allow-run-as-root $PY rag/svr/task_executor.py ; done + while [ 1 -eq 1 ];do + $PY rag/svr/task_executor.py $1 $2; + done } function watch_broker(){ @@ -29,7 +28,12 @@ function task_bro(){ } task_bro & -task_exe & + +WS=8 +for ((i=0;i mx]))[:4] @@ -237,7 +237,7 @@ class Dealer: if i not in cites: continue for c in cites[i]: assert int(c) < len(chunk_v) - res += "##%s$$" % "$".join(cites[i]) + for c in cites[i]: res += f" ##{c}$$" return res diff --git a/rag/nlp/term_weight.py b/rag/nlp/term_weight.py index 3c60808fd..c17cbb072 100644 --- a/rag/nlp/term_weight.py +++ b/rag/nlp/term_weight.py @@ -152,6 +152,7 @@ class Dealer: def ner(t): if not self.ne or t not in self.ne: return 1 + if re.match(r"[0-9,.]+$", t): return 2 m = {"toxic": 2, "func": 1, "corp": 3, "loca": 3, "sch": 3, "stock": 3, "firstnm": 1} return m[self.ne[t]] diff --git a/rag/settings.py b/rag/settings.py index 3c13eefaf..2906ea356 100644 --- a/rag/settings.py +++ b/rag/settings.py @@ -36,3 +36,5 @@ es_logger = getLogger("es") minio_logger = getLogger("minio") cron_logger = getLogger("cron_logger") chunk_logger = getLogger("chunk_logger") +database_logger = getLogger("database") + diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 43fa86474..9ba17eff4 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -23,13 +23,14 @@ import re import sys import traceback from functools import partial -from timeit import default_timer as timer + +from rag.settings import database_logger +from rag.settings import cron_logger, DOC_MAXIMUM_SIZE import numpy as np from elasticsearch_dsl import Q from api.db.services.task_service import TaskService -from rag.settings import cron_logger, DOC_MAXIMUM_SIZE from rag.utils import ELASTICSEARCH from rag.utils import MINIO from rag.utils import rmSpace, findMaxTm @@ -43,7 +44,6 @@ from rag.app import laws, paper, presentation, manual, qa, table, book, resume, from api.db import LLMType, ParserType from api.db.services.document_service import DocumentService from api.db.services.llm_service import LLMBundle -from api.settings import database_logger from api.utils.file_utils import get_project_base_directory BATCH_SIZE = 64 @@ -267,4 +267,4 @@ if __name__ == "__main__": from mpi4py import MPI comm = MPI.COMM_WORLD - main(comm.Get_size(), comm.Get_rank()) + main(int(sys.argv[2]), int(sys.argv[1]))