codingutf8 2022 the open ai team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tokenization classes for openai jukebox import json import os import re import unicodedata from json encoder import infinity from typing import any dict list optional tuple union import numpy as np import regex from tokenizationutils import addedtoken pretrainedtokenizer from tokenizationutilsbase import batchencoding from utils import tensortype isflaxavailable istfavailable istorchavailable logging from utils generic import isjax isnumpy logger logging getloggername vocabfilesnames artistsfile artists json lyricsfile lyrics json genresfile genres json pretrainedvocabfilesmap artistsfile jukebox https huggingface coarthurzjukeboxblobmainartists json genresfile jukebox https huggingface coarthurzjukeboxblobmaingenres json lyricsfile jukebox https huggingface coarthurzjukeboxblobmainlyrics json pretrainedlyrictokenssizes jukebox 512 class jukeboxtokenizerpretrainedtokenizer vocabfilesnames vocabfilesnames pretrainedvocabfilesmap pretrainedvocabfilesmap maxlyricinputsize pretrainedlyrictokenssizes modelinputnames inputids attentionmask def init self artistsfile genresfile lyricsfile versionv3 v2 v2 maxnlyrictokens512 ngenres5 unktokenendoftext kwargs unktoken addedtokenunktoken lstripfalse rstripfalse if isinstanceunktoken str else unktoken self version version self maxnlyrictokens maxnlyrictokens self ngenres ngenres self addedtokensdecoder 0 unktoken with openartistsfile encodingutf8 as vocabhandle self artistsencoder json loadvocabhandle with opengenresfile encodingutf8 as vocabhandle self genresencoder json loadvocabhandle with openlyricsfile encodingutf8 as vocabhandle self lyricsencoder json loadvocabhandle oov razaz09 tn in v2 we had a nvocab80 and in v3 we missed and so nvocab79 of characters if lenself lyricsencoder 79 oov oov replacer r self outofvocab regex compileoov self artistsdecoder v k for k v in self artistsencoder items self genresdecoder v k for k v in self genresencoder items self lyricsdecoder v k for k v in self lyricsencoder items super init unktokenunktoken ngenresngenres versionversion maxnlyrictokensmaxnlyrictokens kwargs property def vocabsizeself return lenself artistsencoder lenself genresencoder lenself lyricsencoder def getvocabself return artistsencoder self artistsencoder genresencoder self genresencoder lyricsencoder self lyricsencoder def converttokentoidself listartists listgenres listlyrics artistsid self artistsencoder getartist 0 for artist in listartists for genres in rangelenlistgenres listgenresgenres self genresencoder getgenre 0 for genre in listgenresgenres listgenresgenres listgenresgenres 1 self ngenres lenlistgenresgenres lyricids self lyricsencoder getcharacter 0 for character in listlyrics0 return artistsid listgenres lyricids def tokenizeself lyrics only lyrics are not tokenized but character based is easily handled return listlyrics def tokenizeself artist genre lyrics kwargs artist genre lyrics self preparefortokenizationartist genre lyrics lyrics self tokenizelyrics return artist genre lyrics def preparefortokenization self artists str genres str lyrics str issplitintowords bool false tuplestr str str dictstr any for idx in rangelenself version if self versionidx v3 artistsidx artistsidx lower genresidx genresidx lower else artistsidx self normalizeartistsidx v2 genresidx self normalizegenre v2 for genre in genresidx split split is for the full dictionary with combined genres if self version0 v2 self outofvocab regex compilerazaz09 tn vocab abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz0123456789 tn self vocab vocabindex index 1 for index in rangelenvocab self vocabunk 0 self nvocab lenvocab 1 self lyricsencoder self vocab self lyricsdecoder v k for k v in self vocab items self lyricsdecoder0 else self outofvocab regex compilerazaz09 tn lyrics self runstripaccentslyrics lyrics lyrics replace n lyrics self outofvocab sub lyrics return artists genres lyrics def runstripaccentsself text normalizes the input text this process is for the genres and the artist args text str artist or genre string to normalize convert the inner content to tensors args tensortype str or utils tensortype optional the type of tensors to use if str should be one of the values of the enum utils tensortype if unset no modification is done prependbatchaxis int optional defaults to false whether or not to add the batch dimension during the conversion convert to tensortype get a function reference for the correct framework do the tensor conversion in batch convert the raw string to a list of token ids args artist str name of the artist genres str list of genres that will be mixed to condition the audio lyrics str optional defaults to lyrics used to condition the generation saves the tokenizer s vocabulary dictionary to the provided savedirectory args savedirectory str a path to the directory where to saved it will be created if it doesn t exist filenameprefix optionalstr optional a prefix to add to the names of the files saved by the tokenizer converts an index integer in a token str using the vocab args artistsindex int index of the artist in its corresponding dictionary genresindex unionlistint int index of the genre in its corresponding dictionary lyricindex listint list of character indices which each correspond to a character coding utf 8 2022 the open ai team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tokenization classes for openai jukebox constructs a jukebox tokenizer jukebox can be conditioned on 3 different inputs artists unique ids are associated to each artist from the provided dictionary genres unique ids are associated to each genre from the provided dictionary lyrics character based tokenization must be initialized with the list of characters that are inside the vocabulary this tokenizer does not require training it should be able to process a different number of inputs as the conditioning of the model can be done on the three different queries if none is provided defaults values will be used depending on the number of genres on which the model should be conditioned n_genres python from transformers import jukeboxtokenizer tokenizer jukeboxtokenizer from_pretrained openai jukebox 1b lyrics tokenizer alan jackson country rock old town road input_ids tensor 0 0 0 6785 546 41 38 30 76 46 41 49 40 76 44 41 27 30 tensor 0 0 0 145 0 tensor 0 0 0 145 0 you can get around that behavior by passing add_prefix_space true when instantiating this tokenizer or when you call it on some text but since the model was not pretrained this way it might yield a decrease in performance tip if nothing is provided the genres and the artist will either be selected randomly or set to none tip this tokenizer inherits from pretrainedtokenizer which contains most of the main methods users should refer to this superclass for more information regarding those methods however the code does not allow that and only supports composing from various genres args artists_file str path to the vocabulary file which contains a mapping between artists and ids the default file supports both v2 and v3 genres_file str path to the vocabulary file which contain a mapping between genres and ids lyrics_file str path to the vocabulary file which contains the accepted characters for the lyrics tokenization version list str optional default to v3 v2 v2 list of the tokenizer versions the 5b lyrics s top level prior model was trained using v3 instead of v2 n_genres int optional defaults to 1 maximum number of genres to use for composition max_n_lyric_tokens int optional defaults to 512 maximum number of lyric tokens to keep unk_token str optional defaults to endoftext the unknown token a token that is not in the vocabulary cannot be converted to an id and is set to be this token instead in v2 we had a n_vocab 80 and in v3 we missed and so n_vocab 79 of characters converts the artist genre and lyrics tokens to their index using the vocabulary the total_length offset and duration have to be provided in order to select relevant lyrics and add padding to the lyrics token sequence converts a string in a sequence of tokens string using the tokenizer split in words for word based vocabulary or sub words for sub word based vocabularies bpe sentencepieces wordpieces do not take care of added tokens only the lyrics are split into character for the character based vocabulary only lyrics are not tokenized but character based is easily handled converts three strings in a 3 sequence of tokens using the tokenizer performs any necessary transformations before tokenization args artist str the artist name to prepare this will mostly lower the string genres str the genre name to prepare this will mostly lower the string lyrics str the lyrics to prepare is_split_into_words bool optional defaults to false whether or not the input is already pre tokenized e g split into words if set to true the tokenizer assumes the input is already split into words for instance by splitting it on whitespace which it will tokenize this is useful for ner or token classification split is for the full dictionary with combined genres strips accents from a piece of text normalizes the input text this process is for the genres and the artist args text str artist or genre string to normalize convert the inner content to tensors args tensor_type str or utils tensortype optional the type of tensors to use if str should be one of the values of the enum utils tensortype if unset no modification is done prepend_batch_axis int optional defaults to false whether or not to add the batch dimension during the conversion convert to tensortype get a function reference for the correct framework noqa f811 do the tensor conversion in batch noqa e722 convert the raw string to a list of token ids args artist str name of the artist genres str list of genres that will be mixed to condition the audio lyrics str optional defaults to lyrics used to condition the generation saves the tokenizer s vocabulary dictionary to the provided save_directory args save_directory str a path to the directory where to saved it will be created if it doesn t exist filename_prefix optional str optional a prefix to add to the names of the files saved by the tokenizer converts an index integer in a token str using the vocab args artists_index int index of the artist in its corresponding dictionary genres_index union list int int index of the genre in its corresponding dictionary lyric_index list int list of character indices which each correspond to a character
import json
import os
import re
import unicodedata
from json.encoder import INFINITY
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import regex

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, is_flax_available, is_tf_available, is_torch_available, logging
from ...utils.generic import _is_jax, _is_numpy


logger = logging.get_logger(__name__)

VOCAB_FILES_NAMES = {
    "artists_file": "artists.json",
    "lyrics_file": "lyrics.json",
    "genres_file": "genres.json",
}

PRETRAINED_VOCAB_FILES_MAP = {
    "artists_file": {
        "jukebox": "https://huggingface.co/ArthurZ/jukebox/blob/main/artists.json",
    },
    "genres_file": {
        "jukebox": "https://huggingface.co/ArthurZ/jukebox/blob/main/genres.json",
    },
    "lyrics_file": {
        "jukebox": "https://huggingface.co/ArthurZ/jukebox/blob/main/lyrics.json",
    },
}

PRETRAINED_LYRIC_TOKENS_SIZES = {
    "jukebox": 512,
}


class JukeboxTokenizer(PreTrainedTokenizer):
    

    vocab_files_names = VOCAB_FILES_NAMES
    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
    max_lyric_input_size = PRETRAINED_LYRIC_TOKENS_SIZES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        artists_file,
        genres_file,
        lyrics_file,
        version=["v3", "v2", "v2"],
        max_n_lyric_tokens=512,
        n_genres=5,
        unk_token="<|endoftext|>",
        **kwargs,
    ):
        unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
        self.version = version
        self.max_n_lyric_tokens = max_n_lyric_tokens
        self.n_genres = n_genres
        self._added_tokens_decoder = {0: unk_token}

        with open(artists_file, encoding="utf-8") as vocab_handle:
            self.artists_encoder = json.load(vocab_handle)

        with open(genres_file, encoding="utf-8") as vocab_handle:
            self.genres_encoder = json.load(vocab_handle)

        with open(lyrics_file, encoding="utf-8") as vocab_handle:
            self.lyrics_encoder = json.load(vocab_handle)

        oov = r"[^A-Za-z0-9.,:;!?\-'\"()\[\] \t\n]+"
        
        if len(self.lyrics_encoder) == 79:
            oov = oov.replace(r"\-'", r"\-+'")

        self.out_of_vocab = regex.compile(oov)
        self.artists_decoder = {v: k for k, v in self.artists_encoder.items()}
        self.genres_decoder = {v: k for k, v in self.genres_encoder.items()}
        self.lyrics_decoder = {v: k for k, v in self.lyrics_encoder.items()}
        super().__init__(
            unk_token=unk_token,
            n_genres=n_genres,
            version=version,
            max_n_lyric_tokens=max_n_lyric_tokens,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return len(self.artists_encoder) + len(self.genres_encoder) + len(self.lyrics_encoder)

    def get_vocab(self):
        return {
            "artists_encoder": self.artists_encoder,
            "genres_encoder": self.genres_encoder,
            "lyrics_encoder": self.lyrics_encoder,
        }

    def _convert_token_to_id(self, list_artists, list_genres, list_lyrics):
        
        artists_id = [self.artists_encoder.get(artist, 0) for artist in list_artists]
        for genres in range(len(list_genres)):
            list_genres[genres] = [self.genres_encoder.get(genre, 0) for genre in list_genres[genres]]
            list_genres[genres] = list_genres[genres] + [-1] * (self.n_genres - len(list_genres[genres]))

        lyric_ids = [[self.lyrics_encoder.get(character, 0) for character in list_lyrics[0]], [], []]
        return artists_id, list_genres, lyric_ids

    def _tokenize(self, lyrics):
        
        
        return list(lyrics)

    def tokenize(self, artist, genre, lyrics, **kwargs):
        
        artist, genre, lyrics = self.prepare_for_tokenization(artist, genre, lyrics)
        lyrics = self._tokenize(lyrics)
        return artist, genre, lyrics

    def prepare_for_tokenization(
        self, artists: str, genres: str, lyrics: str, is_split_into_words: bool = False
    ) -> Tuple[str, str, str, Dict[str, Any]]:
        
        for idx in range(len(self.version)):
            if self.version[idx] == "v3":
                artists[idx] = artists[idx].lower()
                genres[idx] = [genres[idx].lower()]
            else:
                artists[idx] = self._normalize(artists[idx]) + ".v2"
                genres[idx] = [
                    self._normalize(genre) + ".v2" for genre in genres[idx].split("_")
                ]  

        if self.version[0] == "v2":
            self.out_of_vocab = regex.compile(r"[^A-Za-z0-9.,:;!?\-'\"()\[\] \t\n]+")
            vocab = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789.,:;!?-+'\"()[] \t\n"
            self.vocab = {vocab[index]: index + 1 for index in range(len(vocab))}
            self.vocab["<unk>"] = 0
            self.n_vocab = len(vocab) + 1
            self.lyrics_encoder = self.vocab
            self.lyrics_decoder = {v: k for k, v in self.vocab.items()}
            self.lyrics_decoder[0] = ""
        else:
            self.out_of_vocab = regex.compile(r"[^A-Za-z0-9.,:;!?\-+'\"()\[\] \t\n]+")

        lyrics = self._run_strip_accents(lyrics)
        lyrics = lyrics.replace("\\", "\n")
        lyrics = self.out_of_vocab.sub("", lyrics), [], []
        return artists, genres, lyrics

    def _run_strip_accents(self, text):
        
        text = unicodedata.normalize("NFD", text)
        output = []
        for char in text:
            cat = unicodedata.category(char)
            if cat == "Mn":
                continue
            output.append(char)
        return "".join(output)

    def _normalize(self, text: str) -> str:
        

        accepted = (
            [chr(i) for i in range(ord("a"), ord("z") + 1)]
            + [chr(i) for i in range(ord("A"), ord("Z") + 1)]
            + [chr(i) for i in range(ord("0"), ord("9") + 1)]
            + ["."]
        )
        accepted = frozenset(accepted)
        pattern = re.compile(r"_+")
        text = "".join([c if c in accepted else "_" for c in text.lower()])
        text = pattern.sub("_", text).strip("_")
        return text

    def convert_lyric_tokens_to_string(self, lyrics: List[str]) -> str:
        return " ".join(lyrics)

    def convert_to_tensors(
        self, inputs, tensor_type: Optional[Union[str, TensorType]] = None, prepend_batch_axis: bool = False
    ):
        
        
        if not isinstance(tensor_type, TensorType):
            tensor_type = TensorType(tensor_type)

        
        if tensor_type == TensorType.TENSORFLOW:
            if not is_tf_available():
                raise ImportError(
                    "Unable to convert output to TensorFlow tensors format, TensorFlow is not installed."
                )
            import tensorflow as tf

            as_tensor = tf.constant
            is_tensor = tf.is_tensor
        elif tensor_type == TensorType.PYTORCH:
            if not is_torch_available():
                raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.")
            import torch

            as_tensor = torch.tensor
            is_tensor = torch.is_tensor
        elif tensor_type == TensorType.JAX:
            if not is_flax_available():
                raise ImportError("Unable to convert output to JAX tensors format, JAX is not installed.")
            import jax.numpy as jnp  

            as_tensor = jnp.array
            is_tensor = _is_jax
        else:
            as_tensor = np.asarray
            is_tensor = _is_numpy

        

        try:
            if prepend_batch_axis:
                inputs = [inputs]

            if not is_tensor(inputs):
                inputs = as_tensor(inputs)
        except:  
            raise ValueError(
                "Unable to create tensor, you should probably activate truncation and/or padding "
                "with 'padding=True' 'truncation=True' to have batched tensors with the same length."
            )

        return inputs

    def __call__(self, artist, genres, lyrics="", return_tensors="pt") -> BatchEncoding:
        
        input_ids = [0, 0, 0]
        artist = [artist] * len(self.version)
        genres = [genres] * len(self.version)

        artists_tokens, genres_tokens, lyrics_tokens = self.tokenize(artist, genres, lyrics)
        artists_id, genres_ids, full_tokens = self._convert_token_to_id(artists_tokens, genres_tokens, lyrics_tokens)

        attention_masks = [-INFINITY] * len(full_tokens[-1])
        input_ids = [
            self.convert_to_tensors(
                [input_ids + [artists_id[i]] + genres_ids[i] + full_tokens[i]], tensor_type=return_tensors
            )
            for i in range(len(self.version))
        ]
        return BatchEncoding({"input_ids": input_ids, "attention_masks": attention_masks})

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        
        if not os.path.isdir(save_directory):
            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
            return

        artists_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["artists_file"]
        )
        with open(artists_file, "w", encoding="utf-8") as f:
            f.write(json.dumps(self.artists_encoder, ensure_ascii=False))

        genres_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["genres_file"]
        )
        with open(genres_file, "w", encoding="utf-8") as f:
            f.write(json.dumps(self.genres_encoder, ensure_ascii=False))

        lyrics_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["lyrics_file"]
        )
        with open(lyrics_file, "w", encoding="utf-8") as f:
            f.write(json.dumps(self.lyrics_encoder, ensure_ascii=False))

        return (artists_file, genres_file, lyrics_file)

    def _convert_id_to_token(self, artists_index, genres_index, lyric_index):
        
        artist = self.artists_decoder.get(artists_index)
        genres = [self.genres_decoder.get(genre) for genre in genres_index]
        lyrics = [self.lyrics_decoder.get(character) for character in lyric_index]
        return artist, genres, lyrics