Fix WordTag decode bug (#1642)

* fix wordtag decode * Update README.md * Update README.md * Update codestyle
PaddlePaddle · Jan 26, 2022 · 3e68984 · 3e68984
1 parent 28734ab
commit 3e68984
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/paddlenlp/taskflow/knowledge_mining.py b/paddlenlp/taskflow/knowledge_mining.py
@@ -191,7 +191,6 @@ class WordTagTask(Task):
     def __init__(self,
                  model,
                  task,
-                 batch_size=1,
                  params_path=None,
                  tag_path=None,
                  term_schema_path=None,
@@ -419,11 +418,12 @@ def _reset_offset(self, pred_words):
     def _decode(self, batch_texts, batch_pred_tags):
         batch_results = []
         for sent_index in range(len(batch_texts)):
+            sent = batch_texts[sent_index]
             tags = [
                 self._index_to_tags[index]
-                for index in batch_pred_tags[sent_index][self.summary_num:-1]
+                for index in batch_pred_tags[sent_index][self.summary_num:len(
+                    sent) + self.summary_num]
             ]
-            sent = batch_texts[sent_index]
             if self._custom:
                 self._custom.parse_customization(sent, tags, prefix=True)
             sent_out = []

diff --git a/paddlenlp/taskflow/named_entity_recognition.py b/paddlenlp/taskflow/named_entity_recognition.py
@@ -87,11 +87,12 @@ def __init__(self, model, task, **kwargs):
     def _decode(self, batch_texts, batch_pred_tags):
         batch_results = []
         for sent_index in range(len(batch_texts)):
+            sent = batch_texts[sent_index]
             tags = [
                 self._index_to_tags[index]
-                for index in batch_pred_tags[sent_index][self.summary_num:-1]
+                for index in batch_pred_tags[sent_index][self.summary_num:len(
+                    sent) + self.summary_num]
             ]
-            sent = batch_texts[sent_index]
             if self._custom:
                 self._custom.parse_customization(sent, tags, prefix=True)
             sent_out = []
@@ -100,12 +101,12 @@ def _decode(self, batch_texts, batch_pred_tags):
             for ind, tag in enumerate(tags):
                 if partial_word == "":
                     partial_word = sent[ind]
-                    tags_out.append(tag.split('-')[1])
+                    tags_out.append(tag.split('-')[-1])
                     continue
                 if tag.startswith("B") or tag.startswith("S") or tag.startswith(
                         "O"):
                     sent_out.append(partial_word)
-                    tags_out.append(tag.split('-')[1])
+                    tags_out.append(tag.split('-')[-1])
                     partial_word = sent[ind]
                     continue
                 partial_word += sent[ind]