Upload ConlluTokenClassificationPipeline

config.json

@@ -8,6 +8,16 @@
     "AutoModel": "modeling_parser.CobaldParser"
   },
   "consecutive_null_limit": 3,
+  "custom_pipelines": {
+    "conllu-parsing": {
+      "impl": "pipeline.ConlluTokenClassificationPipeline",
+      "pt": [
+        "AutoModel"
+      ],
+      "tf": [],
+      "type": "text"
+    }
+  },
   "deepslot_classifier_hidden_size": 256,
   "dependency_classifier_hidden_size": 128,
   "dropout": 0.1,

pipeline.py

@@ -1,3 +1,5 @@
+from typing import override
+
 from transformers import Pipeline
 
 from src.lemmatize_helper import reconstruct_lemma
@@ -15,6 +17,7 @@ class ConlluTokenClassificationPipeline(Pipeline):
         self.tokenizer = tokenizer
         self.sentenizer = sentenizer
 
+    @override
     def _sanitize_parameters(self, output_format: str = 'list', **kwargs):
         if output_format not in ['list', 'str']:
             raise ValueError(
@@ -23,6 +26,7 @@ class ConlluTokenClassificationPipeline(Pipeline):
         # capture output_format for postprocessing
         return {}, {}, {'output_format': output_format}
 
+    @override
     def preprocess(self, inputs: str) -> dict:
         if not isinstance(inputs, str):
             raise ValueError("pipeline input must be string (text)")
@@ -36,9 +40,11 @@ class ConlluTokenClassificationPipeline(Pipeline):
         self._texts = sentences
         return {"words": words}
 
+    @override
     def _forward(self, model_inputs: dict) -> dict:
         return self.model(**model_inputs, inference_mode=True)
 
+    @override
     def postprocess(self, model_outputs: dict, output_format: str) -> list[dict] | str:
         sentences = self._decode_model_output(model_outputs)
         # Format sentences into CoNLL-U string if requested.