




import os

import jsonlines
import numpy as np
from tqdm import tqdm


a_   		=     2048
a_   		=     4096
a_   		=     42
a_   		=     os.environ.pop("PROCESS_TRAIN", "false")
a_   		=     {"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}




def 					__SCREAMING_SNAKE_CASE   (  lowercase_      )       ->	str:

    """simple docstring"""


    def choose_first(lowercase_					,     lowercase_=False      ):
        assert isinstance(lowercase_					,     lowercase_      )
        if len(lowercase_      ) == 1:
            __UpperCamelCase  		=						answer[0]
            return {k: [answer[k]] for k in answer} if is_long_answer else answer
        for a in answer:
            if is_long_answer:
                __UpperCamelCase  		=						{k: [a[k]] for k in a}
            if len(a['''start_token''']      ) > 0:
                break
        return a

    __UpperCamelCase  		=						{'''id''': example['''id''']}
    __UpperCamelCase  		=						example['''annotations''']
    __UpperCamelCase  		=						annotation['''yes_no_answer''']
    if 0 in yes_no_answer or 1 in yes_no_answer:
        __UpperCamelCase  		=						['''yes'''] if 1 in yes_no_answer else ['''no''']
        __UpperCamelCase  		=						__UpperCamelCase  		=						[]
        __UpperCamelCase  		=						__UpperCamelCase  		=						[]
        __UpperCamelCase  		=						['''<cls>''']
    else:
        __UpperCamelCase  		=						['''short''']
        __UpperCamelCase  		=						choose_first(annotation['''short_answers''']      )
        if len(out['''start_token''']      ) == 0:
            # answer will be long if short is not available
            __UpperCamelCase  		=						['''long''']
            __UpperCamelCase  		=						choose_first(annotation['''long_answer''']					,     is_long_answer=lowercase_      )
            __UpperCamelCase  		=						[]
        answer.update(lowercase_      )

    # disregard some samples
    if len(answer['''start_token''']      ) > 1 or answer["start_token"] == answer["end_token"]:
        __UpperCamelCase  		=						True
    else:
        __UpperCamelCase  		=						False

    __UpperCamelCase  		=						['''start_token''', '''end_token''', '''start_byte''', '''end_byte''', '''text''']
    if not all(isinstance(answer[k]					,     lowercase_      ) for k in cols      ):
        raise ValueError('''Issue in ID'''					,     example['''id''']      )

    return answer




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_=False      )       ->	str:

    """simple docstring"""


    __UpperCamelCase  		=						_get_single_answer(lowercase_      )
    # bytes are of no use
    del answer["start_byte"]
    del answer["end_byte"]

    # handle yes_no answers explicitly
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        __UpperCamelCase  		=						example['''document''']['''tokens''']
        __UpperCamelCase  		=						[]
        for i in range(len(doc['''token''']      )      ):
            if not doc["is_html"][i]:
                context.append(doc['''token'''][i]      )
        return {
            "context": " ".join(lowercase_      ),
            "answer": {
                "start_token": -1_00,  # ignore index in cross-entropy
                "end_token": -1_00,  # ignore index in cross-entropy
                "category": answer["category"],
                "span": answer["category"],  # extra
            },
        }

    # later, help in removing all no answers
    if answer["start_token"] == [-1]:
        return {
            "context": "None",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": "null",
                "span": "None",  # extra
            },
        }

    # handling normal samples

    __UpperCamelCase  		=						['''start_token''', '''end_token''']
    answer.update({k: answer[k][0] if len(answer[k]      ) > 0 else answer[k] for k in cols}      )  # e.g. [10] == 10

    __UpperCamelCase  		=						example['''document''']['''tokens''']
    __UpperCamelCase  		=						answer['''start_token''']
    __UpperCamelCase  		=						answer['''end_token''']

    __UpperCamelCase  		=						[]
    for i in range(len(doc['''token''']      )      ):
        if not doc["is_html"][i]:
            context.append(doc['''token'''][i]      )
        else:
            if answer["start_token"] > i:
                start_token -= 1
            if answer["end_token"] > i:
                end_token -= 1
    __UpperCamelCase  		=						''' '''.join(context[start_token:end_token]      )

    # checking above code
    if assertion:
        __UpperCamelCase  		=						doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
        __UpperCamelCase  		=						doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
        __UpperCamelCase  		=						''' '''.join([old[i] for i in range(len(lowercase_      )      ) if not is_html[i]]      )
        if new != old:
            print('''ID:'''					,     example['''id''']      )
            print('''New:'''					,     lowercase_					,     end='''\n'''      )
            print('''Old:'''					,     lowercase_					,     end='''\n\n'''      )

    return {
        "context": " ".join(lowercase_      ),
        "answer": {
            "start_token": start_token,
            "end_token": end_token - 1,  # this makes it inclusive
            "category": answer["category"],  # either long or short
            "span": new,  # extra
        },
    }




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_					,     lowercase_=20_48					,     lowercase_=40_96					,     lowercase_=True      )       ->	List[str]:

    """simple docstring"""


    __UpperCamelCase  		=						get_context_and_ans(lowercase_					,     assertion=lowercase_      )
    __UpperCamelCase  		=						out['''answer''']

    # later, removing these samples
    if answer["start_token"] == -1:
        return {
            "example_id": example["id"],
            "input_ids": [[-1]],
            "labels": {
                "start_token": [-1],
                "end_token": [-1],
                "category": ["null"],
            },
        }

    __UpperCamelCase  		=						tokenizer(example['''question''']['''text''']					,     out['''context''']      ).input_ids
    __UpperCamelCase  		=						input_ids.index(tokenizer.sep_token_id      ) + 1

    # return yes/no
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        __UpperCamelCase  		=						[]
        __UpperCamelCase  		=						[]
        __UpperCamelCase  		=						input_ids[:q_len]
        __UpperCamelCase  		=						range(lowercase_					,     len(lowercase_      )					,     max_length - doc_stride      )
        for i in doc_start_indices:
            __UpperCamelCase  		=						i + max_length - q_len
            __UpperCamelCase  		=						input_ids[i:end_index]
            inputs.append(q_indices + slice      )
            category.append(answer['''category'''][0]      )
            if slice[-1] == tokenizer.sep_token_id:
                break

        return {
            "example_id": example["id"],
            "input_ids": inputs,
            "labels": {
                "start_token": [-1_00] * len(lowercase_      ),
                "end_token": [-1_00] * len(lowercase_      ),
                "category": category,
            },
        }

    __UpperCamelCase  		=						out['''context'''].split()
    __UpperCamelCase  		=						splitted_context[answer['''end_token''']]
    __UpperCamelCase  		=						len(
        tokenizer(
            ''' '''.join(splitted_context[: answer['''start_token''']]      )					,     add_special_tokens=lowercase_					,     ).input_ids      )
    __UpperCamelCase  		=						len(
        tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]      )					,     add_special_tokens=lowercase_      ).input_ids      )

    answer["start_token"] += q_len
    answer["end_token"] += q_len

    # fixing end token
    __UpperCamelCase  		=						len(tokenizer(lowercase_					,     add_special_tokens=lowercase_      ).input_ids      )
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    __UpperCamelCase  		=						input_ids[answer['''start_token'''] : answer['''end_token'''] + 1]  # right & left are inclusive
    __UpperCamelCase  		=						answer['''start_token''']
    __UpperCamelCase  		=						answer['''end_token''']

    if assertion:
        __UpperCamelCase  		=						tokenizer.decode(lowercase_      )
        if answer["span"] != new:
            print('''ISSUE IN TOKENIZATION'''      )
            print('''OLD:'''					,     answer['''span''']      )
            print('''NEW:'''					,     lowercase_					,     end='''\n\n'''      )

    if len(lowercase_      ) <= max_length:
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
                "category": answer["category"],
            },
        }

    __UpperCamelCase  		=						input_ids[:q_len]
    __UpperCamelCase  		=						range(lowercase_					,     len(lowercase_      )					,     max_length - doc_stride      )

    __UpperCamelCase  		=						[]
    __UpperCamelCase  		=						[]
    __UpperCamelCase  		=						[]
    __UpperCamelCase  		=						[]  # null, yes, no, long, short
    for i in doc_start_indices:
        __UpperCamelCase  		=						i + max_length - q_len
        __UpperCamelCase  		=						input_ids[i:end_index]
        inputs.append(q_indices + slice      )
        assert len(inputs[-1]      ) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            __UpperCamelCase  		=						start_token - i + q_len
            __UpperCamelCase  		=						end_token - i + q_len
            answers_category.append(answer['''category'''][0]      )  # ["short"] -> "short"
        else:
            __UpperCamelCase  		=						-1_00
            __UpperCamelCase  		=						-1_00
            answers_category.append('''null'''      )
        __UpperCamelCase  		=						inputs[-1][start_token : end_token + 1]

        answers_start_token.append(lowercase_      )
        answers_end_token.append(lowercase_      )
        if assertion:
            if new != old and new != [tokenizer.cls_token_id]:
                print('''ISSUE in strided for ID:'''					,     example['''id''']      )
                print('''New:'''					,     tokenizer.decode(lowercase_      )      )
                print('''Old:'''					,     tokenizer.decode(lowercase_      )					,     end='''\n\n'''      )
        if slice[-1] == tokenizer.sep_token_id:
            break

    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        },
    }




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_					,     lowercase_=20_48					,     lowercase_=40_96					,     lowercase_=False      )       ->	str:

    """simple docstring"""


    __UpperCamelCase  		=						get_strided_contexts_and_ans(
        lowercase_					,     lowercase_					,     doc_stride=lowercase_					,     max_length=lowercase_					,     assertion=lowercase_					,     )

    return example




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_      )       ->	Optional[Any]:

    """simple docstring"""


    with jsonlines.open(lowercase_					,     '''a'''      ) as writer:
        for example in tqdm(lowercase_					,     total=len(lowercase_      )					,     desc='''Saving samples ... '''      ):
            __UpperCamelCase  		=						example['''labels''']
            for ids, start, end, cat in zip(
                example['''input_ids''']					,     labels['''start_token''']					,     labels['''end_token''']					,     labels['''category''']					,     ):
                if start == -1 and end == -1:
                    continue  # leave waste samples with no answer
                if cat == "null" and np.random.rand() < 0.6:
                    continue  # removing 50 % samples
                writer.write(
                    {
                        '''input_ids''': ids,
                        '''start_token''': start,
                        '''end_token''': end,
                        '''category''': CATEGORY_MAPPING[cat],
                    }      )


if __name__ == "__main__":
 from datasets import load_dataset

 from transformers import BigBirdTokenizer

 a_   		=     load_dataset("natural_questions")
 a_   		=     BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

 a_   		=     data["train" if PROCESS_TRAIN == "true" else "validation"]

 a_   		=     {
     "tokenizer": tokenizer,
     "doc_stride": DOC_STRIDE,
     "max_length": MAX_LENGTH,
     "assertion": False,
 }
 a_   		=     data.map(prepare_inputs, fn_kwargs=fn_kwargs)
 a_   		=     data.remove_columns(["annotations", "document", "id", "question"])
 print(data)

 np.random.seed(SEED)
 a_   		=     "nq-training.jsonl" if PROCESS_TRAIN == "true" else "nq-validation.jsonl"
 save_to_disk(data, file_name=cache_file_name)












import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
a_   		=     logging.get_logger(__name__)


a_   		=     ["model.decoder.embed_positions.weights"]




def 					__SCREAMING_SNAKE_CASE   (  lowercase_      )       ->	Optional[Any]:

    """simple docstring"""


    if "emb" in name:
        __UpperCamelCase  		=						name.replace('''emb'''					,     '''model.decoder.embed_tokens'''      )
    if "transformer" in name:
        __UpperCamelCase  		=						name.replace('''transformer'''					,     '''model.decoder'''      )
    if "cross_attention" in name:
        __UpperCamelCase  		=						name.replace('''cross_attention'''					,     '''encoder_attn'''      )
    if "linear1" in name:
        __UpperCamelCase  		=						name.replace('''linear1'''					,     '''fc1'''      )
    if "linear2" in name:
        __UpperCamelCase  		=						name.replace('''linear2'''					,     '''fc2'''      )
    if "norm1" in name:
        __UpperCamelCase  		=						name.replace('''norm1'''					,     '''self_attn_layer_norm'''      )
    if "norm_cross" in name:
        __UpperCamelCase  		=						name.replace('''norm_cross'''					,     '''encoder_attn_layer_norm'''      )
    if "norm2" in name:
        __UpperCamelCase  		=						name.replace('''norm2'''					,     '''final_layer_norm'''      )
    if "out_norm" in name:
        __UpperCamelCase  		=						name.replace('''out_norm'''					,     '''model.decoder.layer_norm'''      )
    if "linears" in name:
        __UpperCamelCase  		=						name.replace('''linears'''					,     '''lm_heads'''      )
    if "condition_provider.conditioners.description.output_proj" in name:
        __UpperCamelCase  		=						name.replace('''condition_provider.conditioners.description.output_proj'''					,     '''enc_to_dec_proj'''      )
    return name




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_      )       ->	Tuple[Dict, Dict]:

    """simple docstring"""


    __UpperCamelCase  		=						list(state_dict.keys()      )
    __UpperCamelCase  		=						{}
    for key in keys:
        __UpperCamelCase  		=						state_dict.pop(lowercase_      )
        __UpperCamelCase  		=						rename_keys(lowercase_      )
        if "in_proj_weight" in key:
            # split fused qkv proj
            __UpperCamelCase  		=						val[:hidden_size, :]
            __UpperCamelCase  		=						val[hidden_size : 2 * hidden_size, :]
            __UpperCamelCase  		=						val[-hidden_size:, :]
        elif "enc_to_dec_proj" in key:
            __UpperCamelCase  		=						val
        else:
            __UpperCamelCase  		=						val
    return state_dict, enc_dec_proj_state_dict




def 					__SCREAMING_SNAKE_CASE   (  lowercase_      )       ->	MusicgenDecoderConfig:

    """simple docstring"""


    if checkpoint == "small":
        # default config values
        __UpperCamelCase  		=						10_24
        __UpperCamelCase  		=						24
        __UpperCamelCase  		=						16
    elif checkpoint == "medium":
        __UpperCamelCase  		=						15_36
        __UpperCamelCase  		=						48
        __UpperCamelCase  		=						24
    elif checkpoint == "large":
        __UpperCamelCase  		=						20_48
        __UpperCamelCase  		=						48
        __UpperCamelCase  		=						32
    else:
        raise ValueError(F"Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."      )
    __UpperCamelCase  		=						MusicgenDecoderConfig(
        hidden_size=lowercase_					,     ffn_dim=hidden_size * 4					,     num_hidden_layers=lowercase_					,     num_attention_heads=lowercase_					,     )
    return config




@torch.no_grad()
def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_=None					,     lowercase_=None					,     lowercase_="cpu"      )       ->	List[Any]:

    """simple docstring"""


    __UpperCamelCase  		=						MusicGen.get_pretrained(lowercase_					,     device=lowercase_      )
    __UpperCamelCase  		=						decoder_config_from_checkpoint(lowercase_      )

    __UpperCamelCase  		=						fairseq_model.lm.state_dict()
    __UpperCamelCase		,      __UpperCamelCase  		=						rename_state_dict(
        lowercase_					,     hidden_size=decoder_config.hidden_size      )

    __UpperCamelCase  		=						TaEncoderModel.from_pretrained('''t5-base'''      )
    __UpperCamelCase  		=						EncodecModel.from_pretrained('''facebook/encodec_32khz'''      )
    __UpperCamelCase  		=						MusicgenForCausalLM(lowercase_      ).eval()

    # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
    __UpperCamelCase		,      __UpperCamelCase  		=						decoder.load_state_dict(lowercase_					,     strict=lowercase_      )

    for key in missing_keys.copy():
        if key.startswith(('''text_encoder''', '''audio_encoder''')      ) or key in EXPECTED_MISSING_KEYS:
            missing_keys.remove(lowercase_      )

    if len(lowercase_      ) > 0:
        raise ValueError(F"Missing key(s) in state_dict: {missing_keys}"      )

    if len(lowercase_      ) > 0:
        raise ValueError(F"Unexpected key(s) in state_dict: {unexpected_keys}"      )

    # init the composite model
    __UpperCamelCase  		=						MusicgenForConditionalGeneration(text_encoder=lowercase_					,     audio_encoder=lowercase_					,     decoder=lowercase_      )

    # load the pre-trained enc-dec projection (from the decoder state dict)
    model.enc_to_dec_proj.load_state_dict(lowercase_      )

    # check we can do a forward pass
    __UpperCamelCase  		=						torch.arange(0					,     8					,     dtype=torch.long      ).reshape(2					,     -1      )
    __UpperCamelCase  		=						input_ids.reshape(2 * 4					,     -1      )

    with torch.no_grad():
        __UpperCamelCase  		=						model(input_ids=lowercase_					,     decoder_input_ids=lowercase_      ).logits

    if logits.shape != (8, 1, 20_48):
        raise ValueError('''Incorrect shape for logits'''      )

    # now construct the processor
    __UpperCamelCase  		=						AutoTokenizer.from_pretrained('''t5-base'''      )
    __UpperCamelCase  		=						AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''					,     padding_side='''left'''      )

    __UpperCamelCase  		=						MusicgenProcessor(feature_extractor=lowercase_					,     tokenizer=lowercase_      )

    # set the appropriate bos/pad token ids
    __UpperCamelCase  		=						20_48
    __UpperCamelCase  		=						20_48

    # set other default generation config params
    __UpperCamelCase  		=						int(30 * audio_encoder.config.frame_rate      )
    __UpperCamelCase  		=						True
    __UpperCamelCase  		=						3.0

    if pytorch_dump_folder is not None:
        Path(lowercase_      ).mkdir(exist_ok=lowercase_      )
        logger.info(F"Saving model {checkpoint} to {pytorch_dump_folder}"      )
        model.save_pretrained(lowercase_      )
        processor.save_pretrained(lowercase_      )

    if repo_id:
        logger.info(F"Pushing model {checkpoint} to {repo_id}"      )
        model.push_to_hub(lowercase_      )
        processor.push_to_hub(lowercase_      )


if __name__ == "__main__":
 a_   		=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--checkpoint",
     default="small",
     type=str,
     help="Checkpoint size of the MusicGen model you'd like to convert. Can be one of: `['small', 'medium', 'large']`.",
 )
 parser.add_argument(
     "--pytorch_dump_folder",
     required=True,
     default=None,
     type=str,
     help="Path to the output PyTorch model directory.",
 )
 parser.add_argument(
     "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
 )
 parser.add_argument(
     "--device", default="cpu", type=str, help="Torch device to run the conversion, either cpu or cuda."
 )

 a_   		=     parser.parse_args()
 convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)






