diff --git a/kazu/conf/TransformersModelForTokenClassificationNerStep/default.yaml b/kazu/conf/TransformersModelForTokenClassificationNerStep/default.yaml
index 4a626f0e..817ad71d 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 509bb5aa..00000000
--- 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"