Timestamp extraction and automagic chunking

src/omnilingual_asr/__init__.py

@@ -38,7 +38,6 @@ def setup_fairseq2_extension(container: DependencyContainer) -> None:
 
 
 def _register_models(container: DependencyContainer) -> None:
-
     # Only adding custom wav2vec2 archs for wav2vec2_ssl model in fs2
     register_omnilingual_asr_wav2vec2_ssl_configs(container)
 

src/omnilingual_asr/datasets/impl/manifest_asr_dataset.py

@@ -60,7 +60,6 @@ def create_reader(
         storage_config: ManifestStorageConfig,
         task_config: AsrTaskConfig,
     ) -> DataReader[Seq2SeqBatch]:
-
         storage = ManifestStorage(
             splits=self.splits,
             manifest_dir=self.manifest_dir,

src/omnilingual_asr/datasets/storage/mixture_parquet_storage.py

@@ -205,7 +205,6 @@ def is_train_streaming(split: str, sync_mode: SyncMode) -> bool:
 
     @override
     def create_raw_data_pipeline(self, split: str, gangs: Gangs) -> DataPipelineBuilder:
-
         config = self.config
         schema: LangASRSchema = LangASRSchema()
 

src/omnilingual_asr/datasets/tasks/asr_task.py

@@ -156,7 +156,6 @@ def apply_processing_pipeline(  # type: ignore[override]
         tokenizer: Tokenizer,
         dtype: torch.dtype,
     ) -> DataPipelineBuilder:
-
         config = self.config
 
         # Filtering audio to optimize before batching
@@ -246,7 +245,6 @@ def add_tokenization_pipeline(
         text_selector: str,
         audio_length_selector: str,
     ) -> DataPipelineBuilder:
-
         builder = filter_empty_text(builder, text_selector=text_selector)
 
         builder = filter_fast_speech(
@@ -262,9 +260,13 @@ def add_tokenization_pipeline(
             text_selector=text_selector,
         )
 
-        builder = filter_unknown_sequences(
-            builder, unk_idx=tokenizer.vocab_info.unk_idx, text_selector=text_selector  # type: ignore
-        )
+        unk_idx = tokenizer.vocab_info.unk_idx
+        if unk_idx is not None:
+            builder = filter_unknown_sequences(
+                builder,
+                unk_idx=unk_idx,
+                text_selector=text_selector,  # type: ignore
+            )
 
         if filter_long_text_threshold is not None:
             builder = filter_long_text(
@@ -273,9 +275,10 @@ def add_tokenization_pipeline(
                 text_selector=text_selector,
             )
 
-        if remove_unknown:
+        if remove_unknown and unk_idx is not None:
             builder = filter_unknown_tokens(
-                builder, unk_idx=tokenizer.vocab_info.unk_idx  # type: ignore
+                builder,
+                unk_idx=unk_idx,  # type: ignore
             )
         return builder
 
@@ -293,7 +296,6 @@ def add_bucketing_pipeline(
         batch_size: int,
         no_padding: bool,
     ) -> DataPipelineBuilder:
-
         if batching is BatchingStrategy.LENGTH:
             builder = add_length_batching(
                 builder,
@@ -358,7 +360,6 @@ def add_audio_processing_pipeline(
         npc: int,
         unified_audio_feature_keys: bool,
     ) -> DataPipelineBuilder:
-
         builder = add_audio_decoding(
             builder,
             dtype=dtype,
@@ -412,7 +413,6 @@ def add_postprocessing_pipeline(
         num_prefetch: int,
         no_padding: bool,
     ) -> DataPipelineBuilder:
-
         if no_padding:
             log.warning(
                 "Collating without padding is currently not supported, defaulting to padding."

src/omnilingual_asr/datasets/tasks/ssl_task.py

@@ -130,7 +130,6 @@ def apply_processing_pipeline(  # type: ignore[override]
         gangs: Gangs,
         dtype: torch.dtype,
     ) -> DataPipelineBuilder:
-
         config = self.config
 
         # Shuffle individual samples
@@ -218,7 +217,6 @@ def add_bucketing_pipeline(
         batch_size: int,
         no_padding: bool,
     ) -> DataPipelineBuilder:
-
         if batching is BatchingStrategy.LENGTH:
             builder = add_length_batching(
                 builder,
@@ -260,7 +258,6 @@ def add_audio_processing_pipeline(
         no_padding: bool,
         seed: int,
     ) -> DataPipelineBuilder:
-
         builder = add_audio_decoding(
             builder,
             dtype=dtype,
@@ -320,7 +317,6 @@ def add_postprocessing_pipeline(
         num_prefetch: int,
         no_padding: bool,
     ) -> DataPipelineBuilder:
-
         collater = Collater(pad_value=None if no_padding else 0)
 
         builder.map(collater, num_parallel_calls=npc)