chage tas execution logic (#103)

rag/nlp/__init__.py

@@ -119,7 +119,6 @@ def add_positions(d, poss):
         d["page_num_int"].append(pn + 1)
         d["top_int"].append(top)
         d["position_int"].append((pn + 1, left, right, top, bottom))
-    d["top_int"] = d["top_int"][:1]
 
 
 def remove_contents_table(sections, eng=False):

rag/nlp/query.py

@@ -157,11 +157,11 @@ class EsQueryer:
         s = 1e-9
         for k, v in qtwt.items():
             if k in dtwt:
-                s += v * dtwt[k]
+                s += v# * dtwt[k]
         q = 1e-9
         for k, v in qtwt.items():
             q += v * v
         d = 1e-9
         for k, v in dtwt.items():
             d += v * v
-        return s / math.sqrt(q) / math.sqrt(d)
+        return s / q#math.sqrt(q) / math.sqrt(d)

rag/nlp/search.py

@@ -192,7 +192,7 @@ class Dealer:
         return [float(t) for t in txt.split("\t")]
 
     def insert_citations(self, answer, chunks, chunk_v,
-                         embd_mdl, tkweight=0.3, vtweight=0.7):
+                         embd_mdl, tkweight=0.7, vtweight=0.3):
         assert len(chunks) == len(chunk_v)
         pieces = re.split(r"([；。？!！\n]|[a-z][.?;!][ \n])", answer)
         for i in range(1, len(pieces)):
@@ -224,7 +224,7 @@ class Dealer:
                                                             chunks_tks,
                                                             tkweight, vtweight)
             mx = np.max(sim) * 0.99
-            if mx < 0.55:
+            if mx < 0.35:
                 continue
             cites[idx[i]] = list(
                 set([str(ii) for ii in range(len(chunk_v)) if sim[ii] > 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
 

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]]

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")
+

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]))