
"""simple docstring"""



import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
 from huggingface_hub import snapshot_download
 from pyctcdecode import BeamSearchDecoderCTC

 from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
 from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
 from transformers import WavaVecaForCTC




@require_pyctcdecode
class 		a_    (			unittest.TestCase						):







    def   _snake_case       (      self      :							Tuple  )		->int:






        '''simple docstring'''






        _UpperCAmelCase        = """| <pad> <unk> <s> </s> a b c d e f g h i j k""".split()
        _UpperCAmelCase        = dict(zip(__UpperCamelCase  ,       range(len(__UpperCamelCase  )  )  )  )

        _UpperCAmelCase        = {
            """unk_token""": """<unk>""",
            """bos_token""": """<s>""",
            """eos_token""": """</s>""",
        }
        _UpperCAmelCase        = {
            """feature_size""": 1,
            """padding_value""": 0.0,
            """sampling_rate""": 1_60_00,
            """return_attention_mask""": False,
            """do_normalize""": True,
        }

        _UpperCAmelCase        = tempfile.mkdtemp()
        _UpperCAmelCase        = os.path.join(self.tmpdirname  ,       VOCAB_FILES_NAMES["""vocab_file"""]  )
        _UpperCAmelCase        = os.path.join(self.tmpdirname  ,       __UpperCamelCase  )
        with open(self.vocab_file  ,       """w"""  ,       encoding="""utf-8"""  ) as fp:
            fp.write(json.dumps(__UpperCamelCase  ) + """\n"""  )

        with open(self.feature_extraction_file  ,       """w"""  ,       encoding="""utf-8"""  ) as fp:
            fp.write(json.dumps(__UpperCamelCase  ) + """\n"""  )

        # load decoder from hub
        _UpperCAmelCase        = """hf-internal-testing/ngram-beam-search-decoder"""







    def   _snake_case       (      self      :							List[Any]  ,       **__UpperCamelCase      :							Dict  )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = self.add_kwargs_tokens_map.copy()
        kwargs.update(__UpperCamelCase  )
        return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname  ,       **__UpperCamelCase  )







    def   _snake_case       (      self      :							Tuple  ,       **__UpperCamelCase      :							List[str]  )		->Optional[int]:






        '''simple docstring'''






        return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname  ,       **__UpperCamelCase  )







    def   _snake_case       (      self      :							List[str]  ,       **__UpperCamelCase      :							int  )		->Union[str, Any]:






        '''simple docstring'''






        return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name  ,       **__UpperCamelCase  )







    def   _snake_case       (      self      :							Optional[Any]  )		->Optional[int]:






        '''simple docstring'''






        shutil.rmtree(self.tmpdirname  )







    def   _snake_case       (      self      :							str  )		->List[str]:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        processor.save_pretrained(self.tmpdirname  )
        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname  )

        # tokenizer
        self.assertEqual(processor.tokenizer.get_vocab()  ,       tokenizer.get_vocab()  )
        self.assertIsInstance(processor.tokenizer  ,       __UpperCamelCase  )

        # feature extractor
        self.assertEqual(processor.feature_extractor.to_json_string()  ,       feature_extractor.to_json_string()  )
        self.assertIsInstance(processor.feature_extractor  ,       __UpperCamelCase  )

        # decoder
        self.assertEqual(processor.decoder._alphabet.labels  ,       decoder._alphabet.labels  )
        self.assertEqual(
            processor.decoder.model_container[decoder._model_key]._unigram_set  ,       decoder.model_container[decoder._model_key]._unigram_set  ,       )
        self.assertIsInstance(processor.decoder  ,       __UpperCamelCase  )







    def   _snake_case       (      self      :							List[str]  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = WavaVecaProcessorWithLM(
            tokenizer=self.get_tokenizer()  ,       feature_extractor=self.get_feature_extractor()  ,       decoder=self.get_decoder()  )
        processor.save_pretrained(self.tmpdirname  )

        # make sure that error is thrown when decoder alphabet doesn't match
        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained(
            self.tmpdirname  ,       alpha=5.0  ,       beta=3.0  ,       score_boundary=-7.0  ,       unk_score_offset=3  )

        # decoder
        self.assertEqual(processor.language_model.alpha  ,       5.0  )
        self.assertEqual(processor.language_model.beta  ,       3.0  )
        self.assertEqual(processor.language_model.score_boundary  ,       -7.0  )
        self.assertEqual(processor.language_model.unk_score_offset  ,       3  )







    def   _snake_case       (      self      :							Optional[Any]  )		->Tuple:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_tokenizer()
        # add token to trigger raise
        tokenizer.add_tokens(["""xx"""]  )
        with self.assertRaisesRegex(__UpperCamelCase  ,       """include"""  ):
            WavaVecaProcessorWithLM(
                tokenizer=__UpperCamelCase  ,       feature_extractor=self.get_feature_extractor()  ,       decoder=self.get_decoder()  )







    def   _snake_case       (      self      :							str  )		->str:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        _UpperCAmelCase        = floats_list((3, 10_00)  )

        _UpperCAmelCase        = feature_extractor(__UpperCamelCase  ,       return_tensors="""np"""  )
        _UpperCAmelCase        = processor(__UpperCamelCase  ,       return_tensors="""np"""  )

        for key in input_feat_extract.keys():
            self.assertAlmostEqual(input_feat_extract[key].sum()  ,       input_processor[key].sum()  ,       delta=1e-2  )







    def   _snake_case       (      self      :							Dict  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        _UpperCAmelCase        = """This is a test string"""

        _UpperCAmelCase        = processor(text=__UpperCamelCase  )

        _UpperCAmelCase        = tokenizer(__UpperCamelCase  )

        for key in encoded_tok.keys():
            self.assertListEqual(encoded_tok[key]  ,       encoded_processor[key]  )







    def   _snake_case       (      self      :							Optional[Any]  ,       __UpperCamelCase      :							Tuple=(2, 10, 16)  ,       __UpperCamelCase      :							Dict=77  )		->Optional[Any]:






        '''simple docstring'''






        np.random.seed(__UpperCamelCase  )
        return np.random.rand(*__UpperCamelCase  )







    def   _snake_case       (      self      :							Dict  )		->int:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        _UpperCAmelCase        = self._get_dummy_logits(shape=(10, 16)  ,       seed=13  )

        _UpperCAmelCase        = processor.decode(__UpperCamelCase  )

        _UpperCAmelCase        = decoder.decode_beams(__UpperCamelCase  )[0]

        self.assertEqual(decoded_decoder[0]  ,       decoded_processor.text  )
        self.assertEqual("""</s> <s> </s>"""  ,       decoded_processor.text  )
        self.assertEqual(decoded_decoder[-2]  ,       decoded_processor.logit_score  )
        self.assertEqual(decoded_decoder[-1]  ,       decoded_processor.lm_score  )







    @parameterized.expand([[None], ["""fork"""], ["""spawn"""]]  )
    def   _snake_case       (      self      :							str  ,       __UpperCamelCase      :							Union[str, Any]  )		->Union[str, Any]:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        _UpperCAmelCase        = self._get_dummy_logits()

        # note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
        #       otherwise, the LM won't be available to the pool's sub-processes.
        # manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
        if pool_context is None:
            _UpperCAmelCase        = processor.batch_decode(__UpperCamelCase  )
        else:
            with get_context(__UpperCamelCase  ).Pool() as pool:
                _UpperCAmelCase        = processor.batch_decode(__UpperCamelCase  ,       __UpperCamelCase  )

        _UpperCAmelCase        = list(__UpperCamelCase  )

        with get_context("""fork"""  ).Pool() as p:
            _UpperCAmelCase        = decoder.decode_beams_batch(__UpperCamelCase  ,       __UpperCamelCase  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = [], [], []
        for beams in decoded_beams:
            texts_decoder.append(beams[0][0]  )
            logit_scores_decoder.append(beams[0][-2]  )
            lm_scores_decoder.append(beams[0][-1]  )

        self.assertListEqual(__UpperCamelCase  ,       decoded_processor.text  )
        self.assertListEqual(["""<s> <s> </s>""", """<s> <s> <s>"""]  ,       decoded_processor.text  )
        self.assertListEqual(__UpperCamelCase  ,       decoded_processor.logit_score  )
        self.assertListEqual(__UpperCamelCase  ,       decoded_processor.lm_score  )







    def   _snake_case       (      self      :							Tuple  )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        _UpperCAmelCase        = self._get_dummy_logits()

        _UpperCAmelCase        = 15
        _UpperCAmelCase        = -2_0.0
        _UpperCAmelCase        = -4.0

        _UpperCAmelCase        = processor.batch_decode(
            __UpperCamelCase  ,       beam_width=__UpperCamelCase  ,       beam_prune_logp=__UpperCamelCase  ,       token_min_logp=__UpperCamelCase  ,       )
        _UpperCAmelCase        = decoded_processor_out.text

        _UpperCAmelCase        = list(__UpperCamelCase  )

        with get_context("""fork"""  ).Pool() as pool:
            _UpperCAmelCase        = decoder.decode_beams_batch(
                __UpperCamelCase  ,       __UpperCamelCase  ,       beam_width=__UpperCamelCase  ,       beam_prune_logp=__UpperCamelCase  ,       token_min_logp=__UpperCamelCase  ,       )

        _UpperCAmelCase        = [d[0][0] for d in decoded_decoder_out]
        _UpperCAmelCase        = [d[0][2] for d in decoded_decoder_out]
        _UpperCAmelCase        = [d[0][3] for d in decoded_decoder_out]

        self.assertListEqual(__UpperCamelCase  ,       __UpperCamelCase  )
        self.assertListEqual(["""</s> <s> <s>""", """<s> <s> <s>"""]  ,       __UpperCamelCase  )

        self.assertTrue(np.array_equal(__UpperCamelCase  ,       decoded_processor_out.logit_score  )  )
        self.assertTrue(np.allclose([-2_0.0_5_4, -1_8.4_4_7]  ,       __UpperCamelCase  ,       atol=1e-3  )  )

        self.assertTrue(np.array_equal(__UpperCamelCase  ,       decoded_processor_out.lm_score  )  )
        self.assertTrue(np.allclose([-1_5.5_5_4, -1_3.9_4_7_4]  ,       __UpperCamelCase  ,       atol=1e-3  )  )







    def   _snake_case       (      self      :							int  )		->int:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        _UpperCAmelCase        = self._get_dummy_logits()

        _UpperCAmelCase        = 2.0
        _UpperCAmelCase        = 5.0
        _UpperCAmelCase        = -2_0.0
        _UpperCAmelCase        = True

        _UpperCAmelCase        = processor.batch_decode(
            __UpperCamelCase  ,       alpha=__UpperCamelCase  ,       beta=__UpperCamelCase  ,       unk_score_offset=__UpperCamelCase  ,       lm_score_boundary=__UpperCamelCase  ,       )
        _UpperCAmelCase        = decoded_processor_out.text

        _UpperCAmelCase        = list(__UpperCamelCase  )
        decoder.reset_params(
            alpha=__UpperCamelCase  ,       beta=__UpperCamelCase  ,       unk_score_offset=__UpperCamelCase  ,       lm_score_boundary=__UpperCamelCase  ,       )

        with get_context("""fork"""  ).Pool() as pool:
            _UpperCAmelCase        = decoder.decode_beams_batch(
                __UpperCamelCase  ,       __UpperCamelCase  ,       )

        _UpperCAmelCase        = [d[0][0] for d in decoded_decoder_out]

        self.assertListEqual(__UpperCamelCase  ,       __UpperCamelCase  )
        self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""]  ,       __UpperCamelCase  )
        _UpperCAmelCase        = processor.decoder.model_container[processor.decoder._model_key]
        self.assertEqual(lm_model.alpha  ,       2.0  )
        self.assertEqual(lm_model.beta  ,       5.0  )
        self.assertEqual(lm_model.unk_score_offset  ,       -2_0.0  )
        self.assertEqual(lm_model.score_boundary  ,       __UpperCamelCase  )







    def   _snake_case       (      self      :							Union[str, Any]  )		->Tuple:






        '''simple docstring'''






        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )

        _UpperCAmelCase        = processor.decoder.model_container[processor.decoder._model_key]
        _UpperCAmelCase        = Path(language_model._kenlm_model.path.decode("""utf-8"""  )  ).parent.parent.absolute()

        _UpperCAmelCase        = os.listdir(__UpperCamelCase  )
        _UpperCAmelCase        = ["""alphabet.json""", """language_model"""]

        downloaded_decoder_files.sort()
        expected_decoder_files.sort()

        # test that only decoder relevant files from
        # https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
        # are downloaded and none of the rest (e.g. README.md, ...)
        self.assertListEqual(__UpperCamelCase  ,       __UpperCamelCase  )







    def   _snake_case       (      self      :							str  )		->str:






        '''simple docstring'''






        _UpperCAmelCase        = snapshot_download("""hf-internal-testing/processor_with_lm"""  )

        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained(__UpperCamelCase  )

        _UpperCAmelCase        = processor.decoder.model_container[processor.decoder._model_key]
        _UpperCAmelCase        = Path(language_model._kenlm_model.path.decode("""utf-8"""  )  ).parent.parent.absolute()

        _UpperCAmelCase        = os.listdir(__UpperCamelCase  )
        _UpperCAmelCase        = os.listdir(__UpperCamelCase  )

        local_decoder_files.sort()
        expected_decoder_files.sort()

        # test that both decoder form hub and local files in cache are the same
        self.assertListEqual(__UpperCamelCase  ,       __UpperCamelCase  )







    def   _snake_case       (      self      :							int  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )
        _UpperCAmelCase        = AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""  )

        _UpperCAmelCase        = floats_list((3, 10_00)  )

        _UpperCAmelCase        = processor_wavaveca(__UpperCamelCase  ,       return_tensors="""np"""  )
        _UpperCAmelCase        = processor_auto(__UpperCamelCase  ,       return_tensors="""np"""  )

        for key in input_wavaveca.keys():
            self.assertAlmostEqual(input_wavaveca[key].sum()  ,       input_auto[key].sum()  ,       delta=1e-2  )

        _UpperCAmelCase        = self._get_dummy_logits()

        _UpperCAmelCase        = processor_wavaveca.batch_decode(__UpperCamelCase  )
        _UpperCAmelCase        = processor_auto.batch_decode(__UpperCamelCase  )

        self.assertListEqual(decoded_wavaveca.text  ,       decoded_auto.text  )







    def   _snake_case       (      self      :							Union[str, Any]  )		->List[str]:






        '''simple docstring'''






        _UpperCAmelCase        = self.get_feature_extractor()
        _UpperCAmelCase        = self.get_tokenizer()
        _UpperCAmelCase        = self.get_decoder()

        _UpperCAmelCase        = WavaVecaProcessorWithLM(tokenizer=__UpperCamelCase  ,       feature_extractor=__UpperCamelCase  ,       decoder=__UpperCamelCase  )

        self.assertListEqual(
            processor.model_input_names  ,       feature_extractor.model_input_names  ,       msg="""`processor` and `feature_extractor` model input names do not match"""  ,       )







    @staticmethod
    def   _snake_case       (      __UpperCamelCase      :							Any  ,       __UpperCamelCase      :							Tuple  )		->str:






        '''simple docstring'''






        _UpperCAmelCase        = [d[key] for d in offsets]
        return retrieved_list







    def   _snake_case       (      self      :							Dict  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )
        _UpperCAmelCase        = self._get_dummy_logits()[0]

        _UpperCAmelCase        = processor.decode(__UpperCamelCase  ,       output_word_offsets=__UpperCamelCase  )
        # check Wav2Vec2CTCTokenizerOutput keys for word
        self.assertEqual(len(outputs.keys()  )  ,       4  )
        self.assertTrue("""text""" in outputs  )
        self.assertTrue("""word_offsets""" in outputs  )
        self.assertTrue(isinstance(__UpperCamelCase  ,       __UpperCamelCase  )  )

        self.assertEqual(""" """.join(self.get_from_offsets(outputs["""word_offsets"""]  ,       """word"""  )  )  ,       outputs.text  )
        self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""]  ,       """word"""  )  ,       ["""<s>""", """<s>""", """</s>"""]  )
        self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""]  ,       """start_offset"""  )  ,       [0, 2, 4]  )
        self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""]  ,       """end_offset"""  )  ,       [1, 3, 5]  )







    def   _snake_case       (      self      :							Optional[Any]  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""  )
        _UpperCAmelCase        = self._get_dummy_logits()

        _UpperCAmelCase        = processor.batch_decode(__UpperCamelCase  ,       output_word_offsets=__UpperCamelCase  )

        # check Wav2Vec2CTCTokenizerOutput keys for word
        self.assertEqual(len(outputs.keys()  )  ,       4  )
        self.assertTrue("""text""" in outputs  )
        self.assertTrue("""word_offsets""" in outputs  )
        self.assertTrue(isinstance(__UpperCamelCase  ,       __UpperCamelCase  )  )

        self.assertListEqual(
            [""" """.join(self.get_from_offsets(__UpperCamelCase  ,       """word"""  )  ) for o in outputs["""word_offsets"""]]  ,       outputs.text  )
        self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""][0]  ,       """word"""  )  ,       ["""<s>""", """<s>""", """</s>"""]  )
        self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""][0]  ,       """start_offset"""  )  ,       [0, 2, 4]  )
        self.assertListEqual(self.get_from_offsets(outputs["""word_offsets"""][0]  ,       """end_offset"""  )  ,       [1, 3, 5]  )







    @slow
    @require_torch
    @require_torchaudio
    def   _snake_case       (      self      :							List[str]  )		->str:






        '''simple docstring'''






        import torch

        _UpperCAmelCase        = load_dataset("""common_voice"""  ,       """en"""  ,       split="""train"""  ,       streaming=__UpperCamelCase  )
        _UpperCAmelCase        = ds.cast_column("""audio"""  ,       datasets.Audio(sampling_rate=1_60_00  )  )
        _UpperCAmelCase        = iter(__UpperCamelCase  )
        _UpperCAmelCase        = next(__UpperCamelCase  )

        _UpperCAmelCase        = AutoProcessor.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm"""  )
        _UpperCAmelCase        = WavaVecaForCTC.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm"""  )

        # compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
        _UpperCAmelCase        = processor(sample["""audio"""]["""array"""]  ,       return_tensors="""pt"""  ).input_values

        with torch.no_grad():
            _UpperCAmelCase        = model(__UpperCamelCase  ).logits.cpu().numpy()

        _UpperCAmelCase        = processor.decode(logits[0]  ,       output_word_offsets=__UpperCamelCase  )

        _UpperCAmelCase        = model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
        _UpperCAmelCase        = [
            {
                """start_time""": d["""start_offset"""] * time_offset,
                """end_time""": d["""end_offset"""] * time_offset,
                """word""": d["""word"""],
            }
            for d in output["""word_offsets"""]
        ]

        _UpperCAmelCase        = """WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL"""

        # output words
        self.assertEqual(""" """.join(self.get_from_offsets(__UpperCamelCase  ,       """word"""  )  )  ,       __UpperCamelCase  )
        self.assertEqual(""" """.join(self.get_from_offsets(__UpperCamelCase  ,       """word"""  )  )  ,       output.text  )

        # output times
        _UpperCAmelCase        = torch.tensor(self.get_from_offsets(__UpperCamelCase  ,       """start_time"""  )  )
        _UpperCAmelCase        = torch.tensor(self.get_from_offsets(__UpperCamelCase  ,       """end_time"""  )  )

        # fmt: off
        _UpperCAmelCase        = torch.tensor([1.4_1_9_9, 1.6_5_9_9, 2.2_5_9_9, 3.0, 3.2_4, 3.5_9_9_9, 3.7_9_9_9, 4.0_9_9_9, 4.2_6, 4.9_4, 5.2_8, 5.6_5_9_9, 5.7_8, 5.9_4, 6.3_2, 6.5_3_9_9, 6.6_5_9_9]  )
        _UpperCAmelCase        = torch.tensor([1.5_3_9_9, 1.8_9_9_9, 2.9, 3.1_6, 3.5_3_9_9, 3.7_2, 4.0_1_9_9, 4.1_7_9_9, 4.7_6, 5.1_5_9_9, 5.5_5_9_9, 5.6_9_9_9, 5.8_6, 6.1_9_9_9, 6.3_8, 6.6_1_9_9, 6.9_4]  )
        # fmt: on

        self.assertTrue(torch.allclose(__UpperCamelCase  ,       __UpperCamelCase  ,       atol=0.0_1  )  )
        self.assertTrue(torch.allclose(__UpperCamelCase  ,       __UpperCamelCase  ,       atol=0.0_1  )  )

"""simple docstring"""



import math




class 		a_    :







    def   _snake_case       (      self      :							List[Any]  ,       __UpperCamelCase      :							list[list[float]]  ,       __UpperCamelCase      :							list[int]  )		->int:






        '''simple docstring'''






        _UpperCAmelCase        = 0.0
        _UpperCAmelCase        = 0.0
        for i in range(len(__UpperCamelCase  )  ):
            da += math.pow((sample[i] - weights[0][i])  ,       2  )
            da += math.pow((sample[i] - weights[1][i])  ,       2  )
            return 0 if da > da else 1
        return 0







    def   _snake_case       (      self      :							Dict  ,       __UpperCamelCase      :							list[list[int | float]]  ,       __UpperCamelCase      :							list[int]  ,       __UpperCamelCase      :							int  ,       __UpperCamelCase      :							float  )		->list[list[int | float]]:






        '''simple docstring'''






        for i in range(len(__UpperCamelCase  )  ):
            weights[j][i] += alpha * (sample[i] - weights[j][i])
        return weights







def 			_UpperCamelCase						(    )     ->	None:





    """simple docstring"""

    _UpperCAmelCase        = [[1, 1, 0, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 1, 1]]

    # weight initialization ( n, C )
    _UpperCAmelCase        = [[0.2, 0.6, 0.5, 0.9], [0.8, 0.4, 0.7, 0.3]]

    # training
    _UpperCAmelCase        = SelfOrganizingMap()
    _UpperCAmelCase        = 3
    _UpperCAmelCase        = 0.5

    for _ in range(_A   ):
        for j in range(len(_A   )   ):
            # training sample
            _UpperCAmelCase        = training_samples[j]

            # Compute the winning vector
            _UpperCAmelCase        = self_organizing_map.get_winner(_A       ,   _A   )

            # Update the winning vector
            _UpperCAmelCase        = self_organizing_map.update(_A       ,   _A       ,   _A       ,   _A   )

    # classify test sample
    _UpperCAmelCase        = [0, 0, 0, 1]
    _UpperCAmelCase        = self_organizing_map.get_winner(_A       ,   _A   )

    # results
    print(F"""Clusters that the test sample belongs to : {winner}"""   )
    print(F"""Weights that have been trained : {weights}"""   )


# running the main() function
if __name__ == "__main__":
 main()