improve tag model?

UNDP-Accelerator-Labs · Aug 16, 2024 · 5ad5312 · 5ad5312
1 parent f278835
commit 5ad5312
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/app/system/autotag/cluster.py b/app/system/autotag/cluster.py
@@ -313,7 +313,7 @@ def tag_doc(
             error_msg = f"{error_msg}\nmissing result for {tid}"
             success = False
             continue
-        keywords = tensor_to_str(result["tags"]).split(",")
+        keywords = tensor_to_str(result["tags"]).strip().split(",")
         if not keywords or (len(keywords) == 1 and not keywords[0]):
             continue
         scores = list(result["scores"].cpu().tolist())

diff --git a/nlpapi/tag_model.py b/nlpapi/tag_model.py
@@ -172,8 +172,8 @@ def prep(text: str) -> str:
         if len(texts) == 1:
             # NOTE: fixing the extract_keywords "autocorrect"
             task_keyword_scores = [task_keyword_scores]
-        for task, keyword_scores in zip(
-                inputs.get_current_tasks(), task_keyword_scores):
+        for task_ix, task in enumerate(inputs.get_current_tasks()):
+            keyword_scores = task_keyword_scores[task_ix]
             keywords: list[str] = []
             scores: list[float] = []
             print(keyword_scores)
@@ -187,7 +187,8 @@ def prep(text: str) -> str:
                 [task],
                 {
                     "tags": state.create_single(
-                        str_to_tensor(",".join(keywords))),
+                        str_to_tensor(
+                            ",".join(keywords) if keywords else " ")),
                     "scores": state.create_single(
                         create_tensor(scores, dtype="float")),
                 })