From df54a2cad7c49edd42c4f48e6e68aca1b9818591 Mon Sep 17 00:00:00 2001
From: Richard Jackson <richard.jackson4@astrazeneca.com>
Date: Wed, 11 Sep 2024 11:36:35 +0100
Subject: [PATCH] updated config for
 TransformersModelForTokenClassificationNerStep. Removed obsolete
 multilabel.yaml

---
 .../default.yaml                              | 40 ++++++++++---------
 .../multilabel.yaml                           | 19 ---------
 2 files changed, 22 insertions(+), 37 deletions(-)
 delete mode 100644 kazu/conf/TransformersModelForTokenClassificationNerStep/multilabel.yaml

diff --git a/kazu/conf/TransformersModelForTokenClassificationNerStep/default.yaml b/kazu/conf/TransformersModelForTokenClassificationNerStep/default.yaml
index 4a626f0ed..817ad71d7 100644
--- a/kazu/conf/TransformersModelForTokenClassificationNerStep/default.yaml
+++ b/kazu/conf/TransformersModelForTokenClassificationNerStep/default.yaml
@@ -3,6 +3,9 @@ path: ${oc.env:KAZU_MODEL_PACK}/tinybern
 batch_size: 4
 stride: 16
 max_sequence_length: 128
+keys_to_use: #bert for token classification doesn't use token_type_ids
+  - input_ids
+  - attention_mask
 entity_splitter:
   _target_: kazu.steps.ner.entity_post_processing.NonContiguousEntitySplitter
   entity_conditions:
@@ -13,21 +16,22 @@ entity_splitter:
     disease:
       - _target_: kazu.steps.ner.entity_post_processing.SplitOnConjunctionPattern
         path: ${SciSpacyPipeline.path}
-detect_subspans: False
-threshold: ~
-labels:
-  - 'B-cell_line'
-  - 'B-cell_type'
-  - 'B-disease'
-  - 'B-drug'
-  - 'B-gene'
-  - 'B-species'
-  - 'I-cell_line'
-  - 'I-cell_type'
-  - 'I-disease'
-  - 'I-drug'
-  - 'I-gene'
-  - 'I-species'
-  - 'O'
-strip_re:
-  gene: "( (gene|protein)s?)+$"
+tokenized_word_processor:
+  _target_: kazu.steps.ner.tokenized_word_processor.TokenizedWordProcessor
+  labels:
+    - 'B-cell_line'
+    - 'B-cell_type'
+    - 'B-disease'
+    - 'B-drug'
+    - 'B-gene'
+    - 'B-species'
+    - 'I-cell_line'
+    - 'I-cell_type'
+    - 'I-disease'
+    - 'I-drug'
+    - 'I-gene'
+    - 'I-species'
+    - 'O'
+  strip_re:
+    gene: "( (gene|protein)s?)+$"
+  use_multilabel: false
diff --git a/kazu/conf/TransformersModelForTokenClassificationNerStep/multilabel.yaml b/kazu/conf/TransformersModelForTokenClassificationNerStep/multilabel.yaml
deleted file mode 100644
index 509bb5aae..000000000
--- a/kazu/conf/TransformersModelForTokenClassificationNerStep/multilabel.yaml
+++ /dev/null
@@ -1,19 +0,0 @@
-defaults:
-  - default
-path: ${oc.env:KAZU_MODEL_PACK}/tinybern2
-detect_subspans: true
-threshold: 0.5
-labels:
-  - "O"
-  - "B-cell_line"
-  - "I-cell_line"
-  - "B-cell_type"
-  - "I-cell_type"
-  - "B-disease"
-  - "I-disease"
-  - "B-drug"
-  - "I-drug"
-  - "B-gene"
-  - "I-gene"
-  - "B-species"
-  - "I-species"