Input:
"""Common preprocessing utilities for working with text data"""
import re
import heapq
import os.path as op
from collections import Counter, OrderedDict, defaultdict

import numpy as np

# This list of English stop words is taken from the "Glasgow Information
# Retrieval Group". The original list can be found at
# http://ir.dcs.gla.ac.uk/resources/linguistic_utils/stop_words
_STOP_WORDS = set(
    (
        "a about above across after afterwards again against all almost alone "
        "along already also although always am among amongst amoungst amount an "
        "and another any anyhow anyone anything anyway anywhere are around as at "
        "back be became because become becomes becoming been before beforehand "
        "behind being below beside besides between beyond bill both bottom but by "
        "call can cannot cant co con could couldnt cry de describe detail do done "
        "down due during each eg eight either eleven else elsewhere empty enough "
        "etc even ever every everyone everything everywhere except few fifteen "
        "fifty fill find fire first five for former formerly forty found four from "
        "front full further get give go had has hasnt have he hence her here "
        "hereafter hereby herein hereupon hers herself him himself his how however "
        "hundred i ie if in inc indeed interest into is it its itself keep last "
        "latter latterly least less ltd made many may me meanwhile might mill mine "
        "more moreover most mostly move much must my myself name namely neither "
        "never nevertheless next nine no nobody none noone nor not nothing now "
        "nowhere of off often on once one only onto or other others otherwise our "
        "ours ourselves out over own part per perhaps please put rather re same see "
        "seem seemed seeming seems serious several she should show side since "
        "sincere six sixty so some somehow someone something sometime sometimes "
        "somewhere still such system take ten than that the their them themselves "
        "then thence there thereafter thereby therefore therein thereupon these "
        "they thick thin third this those though three through throughout thru thus "
        "to together too top toward towards twelve twenty two un under until up "
        "upon us very via was we well were what whatever when whence whenever where "
        "whereafter whereas whereby wherein whereupon wherever whether which while "
        "whither who whoever whole whom whose why will with within without would "
        "yet you your yours yourself yourselves"
    ).split(" "),
)

_WORD_REGEX = re.compile(r"(?u)\b\w\w+\b")  # sklearn default
_WORD_REGEX_W_PUNC = re.compile(r"(?u)\w+|[^a-zA-Z0-9\s]")
_WORD_REGEX_W_PUNC_AND_WHITESPACE = re.compile(r"(?u)s?\w+\s?|\s?[^a-zA-Z0-9\s]\s?")

_PUNC_BYTE_REGEX = re.compile(
    r"(33|34|35|36|37|38|39|40|41|42|43|44|45|"
    r"46|47|58|59|60|61|62|63|64|91|92|93|94|"
    r"95|96|123|124|125|126)",
)
_PUNCTUATION = "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
_PUNC_TABLE = str.maketrans("", "", _PUNCTUATION)

def ngrams(sequence, N):
    """Return all `N`-grams of the elements in `sequence`"""
    assert N >= 1
    return list(zip(*[sequence[i:] for i in range(N)]))

def tokenize_whitespace(
    line, lowercase=True, filter_stopwords=True, filter_punctuation=True, **kwargs,
):
    """
    Split a string at any whitespace characters, optionally removing
    punctuation and stop-words in the process.
    """
    line = line.lower() if lowercase else line
    words = line.split()
    line = [strip_punctuation(w) for w in words] if filter_punctuation else line
    return remove_stop_words(words) if filter_stopwords else words

def tokenize_words(
    line, lowercase=True, filter_stopwords=True, filter_punctuation=True, **kwargs,
):
    """
    Split a string into individual words, optionally removing punctuation and
    stop-words in the process.
    """
    REGEX = _WORD_REGEX if filter_punctuation else _WORD_REGEX_W_PUNC
    words = REGEX.findall(line.lower() if lowercase else line)
    return remove_stop_words(words) if filter_stopwords else words

def tokenize_words_bytes(
    line,
    lowercase=True,
    filter_stopwords=True,
    filter_punctuation=True,
    encoding="utf-8",
    **kwargs,
):
    """
    Split a string into individual words, optionally removing punctuation and
    stop-words in the process. Translate each word into a list of bytes.
    """
    words = tokenize_words(
        line,
        lowercase=lowercase,
        filter_stopwords=filter_stopwords,
        filter_punctuation=filter_punctuation,
        **kwargs,
    )
    words = [" ".join([str(i) for i in w.encode(encoding)]) for w in words]
    return words

def tokenize_bytes_raw(line, encoding="utf-8", splitter=None, **kwargs):
    """
    Convert the characters in `line` to a collection of bytes. Each byte is
    represented in decimal as an integer between 0 and 255.

    Parameters
    ----------
    line : str
        The string to tokenize.
    encoding : str
        The encoding scheme for the characters in `line`. Default is `'utf-8'`.
    splitter : {'punctuation', None}
        If `'punctuation'`, split the string at any punctuation character
        before encoding into bytes. If None, do not split `line` at all.
        Default is None.

    Returns
    -------
    bytes : list
        A list of the byte-encoded characters in `line`. Each item in the list
        is a string of space-separated integers between 0 and 255 representing
        the bytes encoding the characters in `line`.
    """
    byte_str = [" ".join([str(i) for i in line.encode(encoding)])]
    if splitter == "punctuation":
        byte_str = _PUNC_BYTE_REGEX.sub(r"-\1-", byte_str[0]).split("-")
    return byte_str

def bytes_to_chars(byte_list, encoding="utf-8"):
    """
    Decode bytes (represented as an integer between 0 and 255) to characters in
    the specified encoding.
    """
    hex_array = [hex(a).replace("0x", "") for a in byte_list]
    hex_array = " ".join([h if len(h) > 1 else f"0{h}" for h in hex_array])
    return bytearray.fromhex(hex_array).decode(encoding)

def tokenize_chars(line, lowercase=True, filter_punctuation=True, **kwargs):
    """
    Split a string into individual characters, optionally removing punctuation
    and stop-words in the process.
    """
    line = line.lower() if lowercase else line
    line = strip_punctuation(line) if filter_punctuation else line
    chars = list(re.sub(" {2,}", " ", line).strip())
    return chars

def remove_stop_words(words):
    """Remove stop words from a list of word strings"""
    return [w for w in words if w.lower() not in _STOP_WORDS]

def strip_punctuation(line):
    """Remove punctuation from a string"""
    return line.translate(_PUNC_TABLE).strip()

#######################################################################
#                          Byte-Pair Encoder                          #
#######################################################################

class BytePairEncoder(object):
    def __init__(self, max_merges=3000, encoding="utf-8"):
        """
        A byte-pair encoder for sub-word embeddings.

        Notes
        -----
        Byte-pair encoding [1][2] is a compression algorithm that iteratively
        replaces the most frequently ocurring byte pairs in a set of documents
        with a new, single token. It has gained popularity as a preprocessing
        step for many NLP tasks due to its simplicity and expressiveness: using
        a base coebook of just 256 unique tokens (bytes), any string can be
        encoded.

        References
        ----------
        .. [1] Gage, P. (1994). A new algorithm for data compression. *C
           Users Journal, 12(2)*, 23–38.
        .. [2] Sennrich, R., Haddow, B., & Birch, A. (2015). Neural machine
           translation of rare words with subword units, *Proceedings of the
           54th Annual Meeting of the Association for Computational
           Linguistics,* 1715-1725.

        Parameters
        ----------
        max_merges : int
            The maximum number of byte pair merges to perform during the
            :meth:`fit` operation. Default is 3000.
        encoding : str
            The encoding scheme for the documents used to train the encoder.
            Default is `'utf-8'`.
        """
        self.parameters = {
            "max_merges": max_merges,
            "encoding": encoding,
        }

        # initialize the byte <-> token and token <-> byte dictionaries. bytes
        # are represented in decimal as integers between 0 and 255. there is a
        # 1:1 correspondence between token and byte representations up to 255.
        self.byte2token = OrderedDict({i: i for i in range(256)})
        self.token2byte = OrderedDict({v: k for k, v in self.byte2token.items()})

    def fit(self, corpus_fps, encoding="utf-8"):
        """
        Train a byte pair codebook on a set of documents.

        Parameters
        ----------
        corpus_fps : str or list of strs
            The filepath / list of filepaths for the document(s) to be used to
            learn the byte pair codebook.
        encoding : str
            The text encoding for documents. Common entries are either 'utf-8'
            (no header byte), or 'utf-8-sig' (header byte). Default is
            'utf-8'.
        """
        vocab = (
            Vocabulary(
                lowercase=False,
                min_count=None,
                max_tokens=None,
                filter_stopwords=False,
                filter_punctuation=False,
                tokenizer="bytes",
            )
            .fit(corpus_fps, encoding=encoding)
            .counts
        )

        # iteratively merge the most common byte bigram across the documents
        for _ in range(self.parameters["max_merges"]):
            pair_counts = self._get_counts(vocab)
            most_common_bigram = max(pair_counts, key=pair_counts.get)
            vocab = self._merge(most_common_bigram, vocab)

        token_bytes = set()
        for k in vocab.keys():
            token_bytes = token_bytes.union([w for w in k.split(" ") if "-" in w])

        for i, t in enumerate(token_bytes):
            byte_tuple = tuple(int(j) for j in t.split("-"))
            self.token2byte[256 + i] = byte_tuple
            self.byte2token[byte_tuple] = 256 + i

        return self

    def _get_counts(self, vocab):
        """Collect bigram counts for the tokens in vocab"""
        pair_counts = defaultdict(int)
        for word, count in vocab.items():
            pairs = ngrams(word.split(" "), 2)
            for p in pairs:
                pair_counts[p] += count
        return pair_counts

    def _merge(self, bigram, vocab):
        """Replace `bigram` with a single token and update vocab accordingly"""
        v_out = {}
        bg = re.escape(" ".join(bigram))
        bigram_regex = re.compile(r"(?<!\S)" + bg + r"(?!\S)")
        for word in vocab.keys():
            # bigram "a b" becomes "a-b"
            w_out = bigram_regex.sub("-".join(bigram), word)
            v_out[w_out] = vocab[word]
        return v_out

    def transform(self, text):
        """
        Transform the words in `text` into their byte pair encoded token IDs.

        Parameters
        ----------
        text: str or list of `N` strings
            The list of strings to encode

        Returns
        -------
        codes : list of `N` lists
            A list of byte pair token IDs for each of the `N` strings in
            `text`.

        Examples
        --------
        >>> B = BytePairEncoder(max_merges=100).fit("./example.txt")
        >>> encoded_tokens = B.transform("Hello! How are you 😁 ?")
        >>> encoded_tokens
        [[72, 879, 474, ...]]
        """
        if isinstance(text, str):
            text = [text]
        return [self._transform(string) for string in text]

    def _transform(self, text):
        """Transform a single text string to a list of byte-pair IDs"""
        P = self.parameters
        _bytes = tokenize_bytes_raw(text, encoding=P["encoding"])

        encoded = []
        for w in _bytes:
            l, r = 0, len(w)
            w = [int(i) for i in w.split(" ")]

            while l < len(w):
                candidate = tuple(w[l:r])

                if len(candidate) > 1 and candidate in self.byte2token:
                    # candidate is a collection of several bytes and is in our
                    # vocab
                    encoded.append(self.byte2token[candidate])
                    l, r = r, len(w)
                elif len(candidate) == 1:
                    # candidate is a single byte and should always be in our
                    # vocab
                    encoded.append(candidate[0])
                    l, r = r, len(w)
                else:
                    # candidate is not in vocab, so we decrease our context
                    # window by 1 and try again
                    r -= 1
        return encoded

    def inverse_transform(self, codes):
        """
        Transform an encoded sequence of byte pair codeword IDs back into
        human-readable text.

        Parameters
        ----------
        codes : list of `N` lists
            A list of `N` lists. Each sublist is a collection of integer
            byte-pair token IDs representing a particular text string.

        Returns
        -------
        text: list of `N` strings
            The decoded strings corresponding to the `N` sublists in `codes`.

        Examples
        --------
        >>> B = BytePairEncoder(max_merges=100).fit("./example.txt")
        >>> encoded_tokens = B.transform("Hello! How are you 😁 ?")
        >>> encoded_tokens
        [[72, 879, 474, ...]]
        >>> B.inverse_transform(encoded_tokens)
        ["Hello! How are you 😁 ?"]
        """
        if isinstance(codes[0], int):
            codes = [codes]

        decoded = []
        P = self.parameters

        for code in codes:
            _bytes = [self.token2byte[t] if t > 255 else [t] for t in code]
            _bytes = [b for blist in _bytes for b in blist]
            decoded.append(bytes_to_chars(_bytes, encoding=P["encoding"]))
        return decoded

    @property
    def codebook(self):
        """
        A list of the learned byte pair codewords, decoded into human-readable
        format
        """
        return [
            self.inverse_transform(t)[0]
            for t in self.byte2token.keys()
            if isinstance(t, tuple)
        ]

    @property
    def tokens(self):
        """A list of the byte pair codeword IDs"""
        return list(self.token2byte.keys())

#######################################################################
#                            Huffman Tree                             #
#######################################################################

class Node(object):
    def __init__(self, key, val):
        self.key = key
        self.val = val
        self.left = None
        self.right = None

    def __gt__(self, other):
        """Greater than"""
        if not isinstance(other, Node):
            return -1
        return self.val > other.val

    def __ge__(self, other):
        """Greater than or equal to"""
        if not isinstance(other, Node):
            return -1
        return self.val >= other.val

    def __lt__(self, other):
        """Less than"""
        if not isinstance(other, Node):
            return -1
        return self.val < other.val

    def __le__(self, other):
        """Less than or equal to"""
        if not isinstance(other, Node):
            return -1
        return self.val <= other.val

class HuffmanEncoder(object):
    def fit(self, text):
        """
        Build a Huffman tree for the tokens in `text` and compute each token's
        binary encoding.

        Notes
        -----
        In a Huffman code, tokens that occur more frequently are (generally)
        represented using fewer bits. Huffman codes produce the minimum expected
        codeword length among all methods for encoding tokens individually.

        Huffman codes correspond to paths through a binary tree, with 1
        corresponding to "move right" and 0 corresponding to "move left". In
        contrast to standard binary trees, the Huffman tree is constructed from the
        bottom up. Construction begins by initializing a min-heap priority queue
        consisting of each token in the corpus, with priority corresponding to the
        token frequency. At each step, the two most infrequent tokens in the corpus
        are removed and become the children of a parent pseudotoken whose
        "frequency" is the sum of the frequencies of its children. This new parent
        pseudotoken is added to the priority queue and the process is repeated
        recursively until no tokens remain.

        Parameters
        ----------
        text: list of strs or :class:`Vocabulary` instance
            The tokenized text or a pretrained :class:`Vocabulary` object to use for
            building the Huffman code.
        """
        self._build_tree(text)
        self._generate_codes()

    def transform(self, text):
        """
        Transform the words in `text` into their Huffman-code representations.

        Parameters
        ----------
        text: list of `N` strings
            The list of words to encode

        Returns
        -------
        codes : list of `N` binary strings
            The encoded words in `text`
        """
        if isinstance(text, str):
            text = [text]
        for token in set(text):
            if token not in self._item2code:
                raise Warning("Token '{}' not in Huffman tree. Skipping".format(token))
        return [self._item2code.get(t, None) for t in text]

    def inverse_transform(self, codes):
        """
        Transform an encoded sequence of bit-strings back into words.

        Parameters
        ----------
        codes : list of `N` binary strings
            A list of encoded bit-strings, represented as strings.

        Returns
        -------
        text: list of `N` strings
            The decoded text.
        """
        if isinstance(codes, str):
            codes = [codes]
        for code in set(codes):
            if code not in self._code2item:
                raise Warning("Code '{}' not in Huffman tree. Skipping".format(code))
        return [self._code2item.get(c, None) for c in codes]

    @property
    def tokens(self):
        """A list the unique tokens in `text`"""
        return list(self._item2code.keys())

    @property
    def codes(self):
        """A list with the Huffman code for each unique token in `text`"""
        return list(self._code2item.keys())

    def _counter(self, text):
        counts = {}
        for item in text:
            counts[item] = counts.get(item, 0) + 1
        return counts

    def _build_tree(self, text):
        """Construct Huffman Tree"""
        PQ = []

        if isinstance(text, Vocabulary):
            counts = text.counts
        else:
            counts = self._counter(text)

        for (k, c) in counts.items():
            PQ.append(Node(k, c))

        # create a priority queue with priority = item frequency
        heapq.heapify(PQ)

        while len(PQ) > 1:
            node1 = heapq.heappop(PQ)  # item with smallest frequency
            node2 = heapq.heappop(PQ)  # item with second smallest frequency

            parent = Node(None, node1.val + node2.val)
            parent.left = node1
            parent.right = node2

            heapq.heappush(PQ, parent)

        self._root = heapq.heappop(PQ)

    def _generate_codes(self):
        current_code = ""
        self._item2code = {}
        self._code2item = {}
        self._build_code(self._root, current_code)

    def _build_code(self, root, current_code):
        if root is None:
            return

        if root.key is not None:
            self._item2code[root.key] = current_code
            self._code2item[current_code] = root.key
            return

        # 0 = move left, 1 = move right
        self._build_code(root.left, current_code + "0")
        self._build_code(root.right, current_code + "1")

#######################################################################
#                             Containers                              #
#######################################################################

class Token:
    def __init__(self, word):
        self.count = 0
        self.word = word

    def __repr__(self):
        """A string representation of the token"""
        return "Token(word='{}', count={})".format(self.word, self.count)

class TFIDFEncoder:
    def __init__(
        self,
        vocab=None,
        lowercase=True,
        min_count=0,
        smooth_idf=True,
        max_tokens=None,
        input_type="files",
        filter_stopwords=True,
        filter_punctuation=True,
        tokenizer="words",
    ):
        r"""
        An object for compiling and encoding the term-frequency
        inverse-document-frequency (TF-IDF) representation of the tokens in a
        text corpus.

        Notes
        -----
        TF-IDF is intended to reflect how important a word is to a document in
        a collection or corpus. For a word token `w` in a document `d`, and a
        corpus, :math:`D = \{d_1, \ldots, d_N\}`, we have:

        .. math::
            \text{TF}(w, d)  &=  \text{num. occurences of }w \text{ in document }d \\
            \text{IDF}(w, D)  &=  \log \frac{|D|}{|\{ d \in D: t \in d \}|}

        Parameters
        ----------
        vocab : :class:`Vocabulary` object or list-like
            An existing vocabulary to filter the tokens in the corpus against.
            Default is None.
        lowercase : bool
            Whether to convert each string to lowercase before tokenization.
            Default is True.
        min_count : int
            Minimum number of times a token must occur in order to be included
            in vocab. Default is 0.
        smooth_idf : bool
            Whether to add 1 to the denominator of the IDF calculation to avoid
            divide-by-zero errors. Default is True.
        max_tokens : int
            Only add the `max_tokens` most frequent tokens that occur more
            than `min_count` to the vocabulary.  If None, add all tokens
            greater that occur more than than `min_count`. Default is None.
        input_type : {'files', 'strings'}
            If 'files', the sequence input to `fit` is expected to be a list
            of filepaths. If 'strings', the input is expected to be a list of
            lists, each sublist containing the raw strings for a single
            document in the corpus. Default is 'filename'.
        filter_stopwords : bool
            Whether to remove stopwords before encoding the words in the
            corpus. Default is True.
        filter_punctuation : bool
            Whether to remove punctuation before encoding the words in the
            corpus. Default is True.
        tokenizer : {'whitespace', 'words', 'characters', 'bytes'}
            Strategy to follow when mapping strings to tokens. The
            `'whitespace'` tokenizer splits strings at whitespace characters.
            The `'words'` tokenizer splits strings using a "word" regex. The
            `'characters'` tokenizer splits strings into individual characters.
            The `'bytes'` tokenizer splits strings into a collection of
            individual bytes.
        """
        # create a function to filter against words in the vocab
        self._filter_vocab = lambda words: words
        if isinstance(vocab, Vocabulary):
            self._filter_vocab = vocab.filter
        elif isinstance(vocab, (list, np.ndarray, set)):
            vocab = set(vocab)
            self._filter_vocab = lambda words: [
                w if w in vocab else "<unk>" for w in words
            ]

        if input_type not in ["files", "strings"]:
            fstr = "`input_type` must be either 'files' or 'strings', but got {}"
            raise ValueError(fstr.format(input_type))

        self._tokens = None
        self._idx2doc = None
        self.term_freq = None
        self.idx2token = None
        self.token2idx = None
        self.inv_doc_freq = None

        self.hyperparameters = {
            "id": "TFIDFEncoder",
            "encoding": None,
            "vocab": vocab
            if not isinstance(vocab, Vocabulary)
            else vocab.hyperparameters,
            "lowercase": lowercase,
            "min_count": min_count,
            "input_type": input_type,
            "max_tokens": max_tokens,
            "smooth_idf": smooth_idf,
            "tokenizer": tokenizer
            if not isinstance(vocab, Vocabulary)
            else vocab.hyperparameters["tokenizer"],
            "filter_stopwords": filter_stopwords
            if not isinstance(vocab, Vocabulary)
            else vocab.hyperparameters["filter_stopwords"],
            "filter_punctuation": filter_punctuation
            if not isinstance(vocab, Vocabulary)
            else vocab.hyperparameters["filter_punctuation"],
        }

    def fit(self, corpus_seq, encoding="utf-8-sig"):
        """
        Compute term-frequencies and inverse document frequencies on a
        collection of documents.

        Parameters
        ----------
        corpus_seq : str or list of strs
            The filepath / list of filepaths / raw string contents of the
            document(s) to be encoded, in accordance with the `input_type`
            parameter passed to the :meth:`__init__` method. Each document is
            expected to be a string of tokens separated by whitespace.
        encoding : str
            Specifies the text encoding for corpus if `input_type` is `files`.
            Common entries are either 'utf-8' (no header byte), or 'utf-8-sig'
            (header byte). Default is 'utf-8-sig'.

        Returns
        -------
        self
        """
        H = self.hyperparameters

        if isinstance(corpus_seq, str):
            corpus_seq = [corpus_seq]

        if H["input_type"] == "files":
            for corpus_fp in corpus_seq:
                assert op.isfile(corpus_fp), "{} does not exist".format(corpus_fp)

        tokens = []
        idx2token, token2idx = {}, {}

        # encode special tokens
        for tt in ["<bol>", "<eol>", "<unk>"]:
            token2idx[tt] = len(tokens)
            idx2token[len(tokens)] = tt
            tokens.append(Token(tt))

        min_count = H["min_count"]
        max_tokens = H["max_tokens"]
        H["encoding"] = encoding

        bol_ix = token2idx["<bol>"]
        eol_ix = token2idx["<eol>"]
        idx2doc, term_freq = {}, {}

        # encode the text in `corpus_fps` without any filtering ...
        for d_ix, doc in enumerate(corpus_seq):
            doc_count = {}
            idx2doc[d_ix] = doc if H["input_type"] == "files" else None
            token2idx, idx2token, tokens, doc_count = self._encode_document(
                doc, token2idx, idx2token, tokens, doc_count, bol_ix, eol_ix,
            )
            term_freq[d_ix] = doc_count

        self._tokens = tokens
        self._idx2doc = idx2doc
        self.token2idx = token2idx
        self.idx2token = idx2token
        self.term_freq = term_freq

        # ... retain only the top `max_tokens` most frequent tokens, coding
        # everything else as <unk> ...
        if max_tokens is not None and len(tokens) > max_tokens:
            self._keep_top_n_tokens()

        # ... replace all words occurring less than `min_count` by <unk> ...
        if min(self._tokens, key=lambda t: t.count).count < min_count:
            self._drop_low_freq_tokens()

        # ... sort tokens alphabetically and reindex ...
        self._sort_tokens()

        # ... finally, calculate inverse document frequency
        self._calc_idf()
        return self

    def _encode_document(
        self, doc, word2idx, idx2word, tokens, doc_count, bol_ix, eol_ix,
    ):
        """Perform tokenization and compute token counts for a single document"""
        H = self.hyperparameters
        lowercase = H["lowercase"]
        filter_stop = H["filter_stopwords"]
        filter_punc = H["filter_punctuation"]

        if H["input_type"] == "files":
            with open(doc, "r", encoding=H["encoding"]) as handle:
                doc = handle.read()

        tokenizer_dict = {
            "words": tokenize_words,
            "characters": tokenize_chars,
            "whitespace": tokenize_whitespace,
            "bytes": tokenize_bytes_raw,
        }
        tokenizer = tokenizer_dict[H["tokenizer"]]

        n_words = 0
        lines = doc.split("\n")
        for line in lines:
            words = tokenizer(
                line,
                lowercase=lowercase,
                filter_stopwords=filter_stop,
                filter_punctuation=filter_punc,
                encoding=H["encoding"],
            )
            words = self._filter_vocab(words)
            n_words += len(words)

            for ww in words:
                if ww not in word2idx:
                    word2idx[ww] = len(tokens)
                    idx2word[len(tokens)] = ww
                    tokens.append(Token(ww))

                t_idx = word2idx[ww]
                tokens[t_idx].count += 1
                doc_count[t_idx] = doc_count.get(t_idx, 0) + 1

            # wrap line in <bol> and <eol> tags
            tokens[bol_ix].count += 1
            tokens[eol_ix].count += 1

            doc_count[bol_ix] = doc_count.get(bol_ix, 0) + 1
            doc_count[eol_ix] = doc_count.get(eol_ix, 0) + 1
        return word2idx, idx2word, tokens, doc_count

    def _keep_top_n_tokens(self):
        N = self.hyperparameters["max_tokens"]
        doc_counts, word2idx, idx2word = {}, {}, {}
        tokens = sorted(self._tokens, key=lambda x: x.count, reverse=True)

        # reindex the top-N tokens...
        unk_ix = None
        for idx, tt in enumerate(tokens[:N]):
            word2idx[tt.word] = idx
            idx2word[idx] = tt.word

            if tt.word == "<unk>":
                unk_ix = idx

        # ... if <unk> isn't in the top-N, add it, replacing the Nth
        # most-frequent word and adjust the <unk> count accordingly ...
        if unk_ix is None:
            unk_ix = self.token2idx["<unk>"]
            old_count = tokens[N - 1].count
            tokens[N - 1] = self._tokens[unk_ix]
            tokens[N - 1].count += old_count
            word2idx["<unk>"] = N - 1
            idx2word[N - 1] = "<unk>"

        # ... and recode all dropped tokens as "<unk>"
        for tt in tokens[N:]:
            tokens[unk_ix].count += tt.count

        # ... finally, reindex the word counts for each document
        doc_counts = {}
        for d_ix in self.term_freq.keys():
            doc_counts[d_ix] = {}
            for old_ix, d_count in self.term_freq[d_ix].items():
                word = self.idx2token[old_ix]
                new_ix = word2idx.get(word, unk_ix)
                doc_counts[d_ix][new_ix] = doc_counts[d_ix].get(new_ix, 0) + d_count

        self._tokens = tokens[:N]
        self.token2idx = word2idx
        self.idx2token = idx2word
        self.term_freq = doc_counts

        assert len(self._tokens) <= N

    def _drop_low_freq_tokens(self):
        """
        Replace all tokens that occur less than `min_count` with the `<unk>`
        token.
        """
        H = self.hyperparameters
        unk_token = self._tokens[self.token2idx["<unk>"]]
        eol_token = self._tokens[self.token2idx["<eol>"]]
        bol_token = self._tokens[self.token2idx["<bol>"]]
        tokens = [unk_token, eol_token, bol_token]

        unk_idx = 0
        word2idx = {"<unk>": 0, "<eol>": 1, "<bol>": 2}
        idx2word = {0: "<unk>", 1: "<eol>", 2: "<bol>"}
        special = {"<eol>", "<bol>", "<unk>"}

        for tt in self._tokens:
            if tt.word not in special:
                if tt.count < H["min_count"]:
                    tokens[unk_idx].count += tt.count
                else:
                    word2idx[tt.word] = len(tokens)
                    idx2word[len(tokens)] = tt.word
                    tokens.append(tt)

        # reindex document counts
        doc_counts = {}
        for d_idx in self.term_freq.keys():
            doc_counts[d_idx] = {}
            for old_idx, d_count in self.term_freq[d_idx].items():
                word = self.idx2token[old_idx]
                new_idx = word2idx.get(word, unk_idx)
                doc_counts[d_idx][new_idx] = doc_counts[d_idx].get(new_idx, 0) + d_count

        self._tokens = tokens
        self.token2idx = word2idx
        self.idx2token = idx2word
        self.term_freq = doc_counts

    def _sort_tokens(self):
        # sort tokens alphabetically and recode
        ix = 0
        token2idx, idx2token, = (
            {},
            {},
        )
        special = ["<eol>", "<bol>", "<unk>"]
        words = sorted(self.token2idx.keys())
        term_freq = {d: {} for d in self.term_freq.keys()}

        for w in words:
            if w not in special:
                old_ix = self.token2idx[w]
                token2idx[w], idx2token[ix] = ix, w
                for d in self.term_freq.keys():
                    if old_ix in self.term_freq[d]:
                        count = self.term_freq[d][old_ix]
                        term_freq[d][ix] = count
                ix += 1

        for w in special:
            token2idx[w] = len(token2idx)
            idx2token[len(idx2token)] = w

        self.token2idx = token2idx
        self.idx2token = idx2token
        self.term_freq = term_freq
        self.vocab_counts = Counter({t.word: t.count for t in self._tokens})

    def _calc_idf(self):
        """
        Compute the (smoothed-) inverse-document frequency for each token in
        the corpus.

        For a word token `w`, the IDF is simply

            IDF(w) = log ( |D| / |{ d in D: w in d }| ) + 1

        where D is the set of all documents in the corpus,

            D = {d1, d2, ..., dD}

        If `smooth_idf` is True, we perform additive smoothing on the number of
        documents containing a given word, equivalent to pretending that there
        exists a final D+1st document that contains every word in the corpus:

            SmoothedIDF(w) = log ( |D| + 1 / [1 + |{ d in D: w in d }|] ) + 1
        """
        inv_doc_freq = {}
        smooth_idf = self.hyperparameters["smooth_idf"]
        tf, doc_idxs = self.term_freq, self._idx2doc.keys()

        D = len(self._idx2doc) + int(smooth_idf)
        for word, w_ix in self.token2idx.items():
            d_count = int(smooth_idf)
            d_count += np.sum([1 if w_ix in tf[d_ix] else 0 for d_ix in doc_idxs])
            inv_doc_freq[w_ix] = 1 if d_count == 0 else np.log(D / d_count) + 1
        self.inv_doc_freq = inv_doc_freq

    def transform(self, ignore_special_chars=True):
        """
        Generate the term-frequency inverse-document-frequency encoding of a
        text corpus.

        Parameters
        ----------
        ignore_special_chars : bool
            Whether to drop columns corresponding to "<eol>", "<bol>", and
            "<unk>" tokens from the final tfidf encoding. Default is True.

        Returns
        -------
        tfidf : numpy array of shape `(D, M [- 3])`
            The encoded corpus, with each row corresponding to a single
            document, and each column corresponding to a token id. The mapping
            between column numbers and tokens is stored in the `idx2token`
            attribute IFF `ignore_special_chars` is False. Otherwise, the
            mappings are not accurate.
        """
        D, N = len(self._idx2doc), len(self._tokens)
        tf = np.zeros((D, N))
        idf = np.zeros((D, N))

        for d_ix in self._idx2doc.keys():
            words, counts = zip(*self.term_freq[d_ix].items())
            docs = np.ones(len(words), dtype=int) * d_ix
            tf[docs, words] = counts

        words = sorted(self.idx2token.keys())
        idf = np.tile(np.array([self.inv_doc_freq[w] for w in words]), (D, 1))
        tfidf = tf * idf

        if ignore_special_chars:
            idxs = [
                self.token2idx["<unk>"],
                self.token2idx["<eol>"],
                self.token2idx["<bol>"],
            ]
            tfidf = np.delete(tfidf, idxs, 1)

        return tfidf

class Vocabulary:
    def __init__(
        self,
        lowercase=True,
        min_count=None,
        max_tokens=None,
        filter_stopwords=True,
        filter_punctuation=True,
        tokenizer="words",
    ):
        """
        An object for compiling and encoding the unique tokens in a text corpus.

        Parameters
        ----------
        lowercase : bool
            Whether to convert each string to lowercase before tokenization.
            Default is True.
        min_count : int
            Minimum number of times a token must occur in order to be included
            in vocab. If `None`, include all tokens from `corpus_fp` in vocab.
            Default is None.
        max_tokens : int
            Only add the `max_tokens` most frequent tokens that occur more
            than `min_count` to the vocabulary.  If None, add all tokens
            that occur more than than `min_count`. Default is None.
        filter_stopwords : bool
            Whether to remove stopwords before encoding the words in the
            corpus. Default is True.
        filter_punctuation : bool
            Whether to remove punctuation before encoding the words in the
            corpus. Default is True.
        tokenizer : {'whitespace', 'words', 'characters', 'bytes'}
            Strategy to follow when mapping strings to tokens. The
            `'whitespace'` tokenizer splits strings at whitespace characters.
            The `'words'` tokenizer splits strings using a "word" regex. The
            `'characters'` tokenizer splits strings into individual characters.
            The `'bytes'` tokenizer splits strings into a collection of
            individual bytes.
        """
        self.hyperparameters = {
            "id": "Vocabulary",
            "encoding": None,
            "corpus_fps": None,
            "lowercase": lowercase,
            "min_count": min_count,
            "max_tokens": max_tokens,
            "filter_stopwords": filter_stopwords,
            "filter_punctuation": filter_punctuation,
            "tokenizer": tokenizer,
        }

    def __len__(self):
        """Return the number of tokens in the vocabulary"""
        return len(self._tokens)

    def __iter__(self):
        """Return an iterator over the tokens in the vocabulary"""
        return iter(self._tokens)

    def __contains__(self, word):
        """Assert whether `word` is a token in the vocabulary"""
        return word in self.token2idx

    def __getitem__(self, key):
        """
        Return the token (if key is an integer) or the index (if key is a string)
        for the key in the vocabulary, if it exists.
        """
        if isinstance(key, str):
            return self._tokens[self.token2idx[key]]
        if isinstance(key, int):
            return self._tokens[key]

    @property
    def n_tokens(self):
        """The number of unique word tokens in the vocabulary"""
        return len(self.token2idx)

    @property
    def n_words(self):
        """The total number of words in the corpus"""
        return sum(self.counts.values())

    @property
    def shape(self):
        """The number of unique word tokens in the vocabulary"""
        return self._tokens.shape

    def most_common(self, n=5):
        """Return the top `n` most common tokens in the corpus"""
        return self.counts.most_common()[:n]

    def words_with_count(self, k):
        """Return all tokens that occur `k` times in the corpus"""
        return [w for w, c in self.counts.items() if c == k]

    def filter(self, words, unk=True):  # noqa: A003
        """
        Filter (or replace) any word in `words` that is not present in
        `Vocabulary`.

        Parameters
        ----------
        words : list of strs
            A list of words to filter
        unk : bool
            Whether to replace any out of vocabulary words in `words` with the
            ``<unk>`` token (True) or skip them entirely (False).  Default is
            True.

        Returns
        -------
        filtered : list of strs
            The list of words filtered against the words in Vocabulary.
        """
        if unk:
            return [w if w in self else "<unk>" for w in words]
        return [w for w in words if w in self]

    def words_to_indices(self, words):
        """
        Convert the words in `words` to their token indices. If a word is not
        in the vocabulary, return the index for the ``<unk>`` token

        Parameters
        ----------
        words : list of strs
            A list of words to filter

        Returns
        -------
        indices : list of ints
            The token indices for each word in `words`
        """
        unk_ix = self.token2idx["<unk>"]
        lowercase = self.hyperparameters["lowercase"]
        words = [w.lower() for w in words] if lowercase else words
        return [self.token2idx[w] if w in self else unk_ix for w in words]

    def indices_to_words(self, indices):
        """
        Convert the indices in `indices` to their word values. If an index is
        not in the vocabulary, return the ``<unk>`` token.

        Parameters
        ----------
        indices : list of ints
            The token indices for each word in `words`

        Returns
        -------
        words : list of strs
            The word strings corresponding to each token index in `indices`
        """
        unk = "<unk>"
        return [self.idx2token[i] if i in self.idx2token else unk for i in indices]

    def fit(self, corpus_fps, encoding="utf-8-sig"):
        """
        Compute the vocabulary across a collection of documents.

        Parameters
        ----------
        corpus_fps : str or list of strs
            The filepath / list of filepaths for the document(s) to be encoded.
            Each document is expected to be encoded as newline-separated
            string of text, with adjacent tokens separated by a whitespace
            character.
        encoding : str
            Specifies the text encoding for corpus. Common entries are either
            'utf-8' (no header byte), or 'utf-8-sig' (header byte). Default is
            'utf-8-sig'.

        Returns
        -------
        self
        """
        if isinstance(corpus_fps, str):
            corpus_fps = [corpus_fps]

        for corpus_fp in corpus_fps:
            assert op.isfile(corpus_fp), "{} does not exist".format(corpus_fp)

        tokens = []
        H = self.hyperparameters
        idx2word, word2idx = {}, {}

        tokenizer_dict = {
            "words": tokenize_words,
            "characters": tokenize_chars,
            "whitespace": tokenize_whitespace,
            "bytes": tokenize_bytes_raw,
        }

        min_count = H["min_count"]
        lowercase = H["lowercase"]
        max_tokens = H["max_tokens"]
        filter_stop = H["filter_stopwords"]
        filter_punc = H["filter_punctuation"]
        tokenizer = tokenizer_dict[H["tokenizer"]]

        H["encoding"] = encoding
        H["corpus_fps"] = corpus_fps

        # encode special tokens
        for tt in ["<bol>", "<eol>", "<unk>"]:
            word2idx[tt] = len(tokens)
            idx2word[len(tokens)] = tt
            tokens.append(Token(tt))

        bol_ix = word2idx["<bol>"]
        eol_ix = word2idx["<eol>"]

        for d_ix, doc_fp in enumerate(corpus_fps):
            with open(doc_fp, "r", encoding=H["encoding"]) as doc:
                for line in doc:
                    words = tokenizer(
                        line,
                        lowercase=lowercase,
                        filter_stopwords=filter_stop,
                        filter_punctuation=filter_punc,
                        encoding=H["encoding"],
                    )

                    for ww in words:
                        if ww not in word2idx:
                            word2idx[ww] = len(tokens)
                            idx2word[len(tokens)] = ww
                            tokens.append(Token(ww))

                        t_idx = word2idx[ww]
                        tokens[t_idx].count += 1

                    # wrap line in <bol> and <eol> tags
                    tokens[bol_ix].count += 1
                    tokens[eol_ix].count += 1

        self._tokens = tokens
        self.token2idx = word2idx
        self.idx2token = idx2word

        # replace all words occurring less than `min_count` by <unk>
        if min_count is not None:
            self._drop_low_freq_tokens()

        # retain only the top `max_tokens` most frequent tokens, coding
        # everything else as <unk>
        if max_tokens is not None and len(tokens) > max_tokens:
            self._keep_top_n_tokens()

        counts = {w: self._tokens[ix].count for w, ix in self.token2idx.items()}
        self.counts = Counter(counts)
        self._tokens = np.array(self._tokens)
        return self

    def _keep_top_n_tokens(self):
        word2idx, idx2word = {}, {}
        N = self.hyperparameters["max_tokens"]
        tokens = sorted(self._tokens, key=lambda x: x.count, reverse=True)

        # reindex the top-N tokens...
        unk_ix = None
        for idx, tt in enumerate(tokens[:N]):
            word2idx[tt.word] = idx
            idx2word[idx] = tt.word

            if tt.word == "<unk>":
                unk_ix = idx

        # ... if <unk> isn't in the top-N, add it, replacing the Nth
        # most-frequent word and adjusting the <unk> count accordingly ...
        if unk_ix is None:
            unk_ix = self.token2idx["<unk>"]
            old_count = tokens[N - 1].count
            tokens[N - 1] = self._tokens[unk_ix]
            tokens[N - 1].count += old_count
            word2idx["<unk>"] = N - 1
            idx2word[N - 1] = "<unk>"

        # ... and recode all dropped tokens as "<unk>"
        for tt in tokens[N:]:
            tokens[unk_ix].count += tt.count

        self._tokens = tokens[:N]
        self.token2idx = word2idx
        self.idx2token = idx2word

        assert len(self._tokens) <= N

    def _drop_low_freq_tokens(self):
        """
        Replace all tokens that occur less than `min_count` with the `<unk>`
        token.
        """
        unk_idx = 0
        unk_token = self._tokens[self.token2idx["<unk>"]]
        eol_token = self._tokens[self.token2idx["<eol>"]]
        bol_token = self._tokens[self.token2idx["<bol>"]]

        H = self.hyperparameters
        tokens = [unk_token, eol_token, bol_token]
        word2idx = {"<unk>": 0, "<eol>": 1, "<bol>": 2}
        idx2word = {0: "<unk>", 1: "<eol>", 2: "<bol>"}
        special = {"<eol>", "<bol>", "<unk>"}

        for tt in self._tokens:
            if tt.word not in special:
                if tt.count < H["min_count"]:
                    tokens[unk_idx].count += tt.count
                else:
                    word2idx[tt.word] = len(tokens)
                    idx2word[len(tokens)] = tt.word
                    tokens.append(tt)

        self._tokens = tokens
        self.token2idx = word2idx
        self.idx2token = idx2word

Output: tokenize_chars

Output: remove_stop_words

Output: strip_punctuation

Output: ngrams

Output: tokenize_whitespace

Output: bytes_to_chars

Output: tokenize_words

Output: tokenize_bytes_raw

Output: tokenize_words_bytes

Output: Token

Output: Node

Output: Vocabulary

Output: BytePairEncoder

Output: TFIDFEncoder

Output: HuffmanEncoder


Input:
"""
Defines helper methods useful for loading and caching Interface examples.
"""
from __future__ import annotations

import ast
import csv
import inspect
import os
import shutil
import subprocess
import tempfile
import warnings
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Iterable, Literal, Optional

import numpy as np
import PIL
import PIL.Image
from gradio_client import utils as client_utils
from gradio_client.documentation import document, set_documentation_group

from gradio import components, oauth, processing_utils, routes, utils, wasm_utils
from gradio.context import Context, LocalContext
from gradio.data_classes import GradioModel, GradioRootModel
from gradio.events import EventData
from gradio.exceptions import Error
from gradio.flagging import CSVLogger

if TYPE_CHECKING:  # Only import for type checking (to avoid circular imports).
    from gradio.components import Component

LOG_FILE = "log.csv"

set_documentation_group("helpers")

def create_examples(
    examples: list[Any] | list[list[Any]] | str,
    inputs: Component | list[Component],
    outputs: Component | list[Component] | None = None,
    fn: Callable | None = None,
    cache_examples: bool = False,
    examples_per_page: int = 10,
    _api_mode: bool = False,
    label: str | None = None,
    elem_id: str | None = None,
    run_on_click: bool = False,
    preprocess: bool = True,
    postprocess: bool = True,
    api_name: str | Literal[False] = "load_example",
    batch: bool = False,
):
    """Top-level synchronous function that creates Examples. Provided for backwards compatibility, i.e. so that gr.Examples(...) can be used to create the Examples component."""
    examples_obj = Examples(
        examples=examples,
        inputs=inputs,
        outputs=outputs,
        fn=fn,
        cache_examples=cache_examples,
        examples_per_page=examples_per_page,
        _api_mode=_api_mode,
        label=label,
        elem_id=elem_id,
        run_on_click=run_on_click,
        preprocess=preprocess,
        postprocess=postprocess,
        api_name=api_name,
        batch=batch,
        _initiated_directly=False,
    )
    examples_obj.create()
    return examples_obj

@document()
class Examples:
    """
    This class is a wrapper over the Dataset component and can be used to create Examples
    for Blocks / Interfaces. Populates the Dataset component with examples and
    assigns event listener so that clicking on an example populates the input/output
    components. Optionally handles example caching for fast inference.

    Demos: blocks_inputs, fake_gan
    Guides: more-on-examples-and-flagging, using-hugging-face-integrations, image-classification-in-pytorch, image-classification-in-tensorflow, image-classification-with-vision-transformers, create-your-own-friends-with-a-gan
    """

    def __init__(
        self,
        examples: list[Any] | list[list[Any]] | str,
        inputs: Component | list[Component],
        outputs: Component | list[Component] | None = None,
        fn: Callable | None = None,
        cache_examples: bool = False,
        examples_per_page: int = 10,
        _api_mode: bool = False,
        label: str | None = "Examples",
        elem_id: str | None = None,
        run_on_click: bool = False,
        preprocess: bool = True,
        postprocess: bool = True,
        api_name: str | Literal[False] = "load_example",
        batch: bool = False,
        _initiated_directly: bool = True,
    ):
        """
        Parameters:
            examples: example inputs that can be clicked to populate specific components. Should be nested list, in which the outer list consists of samples and each inner list consists of an input corresponding to each input component. A string path to a directory of examples can also be provided but it should be within the directory with the python file running the gradio app. If there are multiple input components and a directory is provided, a log.csv file must be present in the directory to link corresponding inputs.
            inputs: the component or list of components corresponding to the examples
            outputs: optionally, provide the component or list of components corresponding to the output of the examples. Required if `cache_examples` is True.
            fn: optionally, provide the function to run to generate the outputs corresponding to the examples. Required if `cache_examples` is True.
            cache_examples: if True, caches examples for fast runtime. If True, then `fn` and `outputs` must be provided. If `fn` is a generator function, then the last yielded value will be used as the output.
            examples_per_page: how many examples to show per page.
            label: the label to use for the examples component (by default, "Examples")
            elem_id: an optional string that is assigned as the id of this component in the HTML DOM.
            run_on_click: if cache_examples is False, clicking on an example does not run the function when an example is clicked. Set this to True to run the function when an example is clicked. Has no effect if cache_examples is True.
            preprocess: if True, preprocesses the example input before running the prediction function and caching the output. Only applies if `cache_examples` is True.
            postprocess: if True, postprocesses the example output after running the prediction function and before caching. Only applies if `cache_examples` is True.
            api_name: Defines how the event associated with clicking on the examples appears in the API docs. Can be a string or False. If set to a string, the endpoint will be exposed in the API docs with the given name. If False, the endpoint will not be exposed in the API docs and downstream apps (including those that `gr.load` this app) will not be able to use the example function.
            batch: If True, then the function should process a batch of inputs, meaning that it should accept a list of input values for each parameter. Used only if cache_examples is True.
        """
        if _initiated_directly:
            warnings.warn(
                "Please use gr.Examples(...) instead of gr.examples.Examples(...) to create the Examples.",
            )

        if cache_examples and (fn is None or outputs is None):
            raise ValueError("If caching examples, `fn` and `outputs` must be provided")

        if not isinstance(inputs, list):
            inputs = [inputs]
        if outputs and not isinstance(outputs, list):
            outputs = [outputs]

        working_directory = Path().absolute()

        if examples is None:
            raise ValueError("The parameter `examples` cannot be None")
        elif isinstance(examples, list) and (
            len(examples) == 0 or isinstance(examples[0], list)
        ):
            pass
        elif (
            isinstance(examples, list) and len(inputs) == 1
        ):  # If there is only one input component, examples can be provided as a regular list instead of a list of lists
            examples = [[e] for e in examples]
        elif isinstance(examples, str):
            if not Path(examples).exists():
                raise FileNotFoundError(
                    f"Could not find examples directory: {examples}"
                )
            working_directory = examples
            if not (Path(examples) / LOG_FILE).exists():
                if len(inputs) == 1:
                    examples = [[e] for e in os.listdir(examples)]
                else:
                    raise FileNotFoundError(
                        "Could not find log file (required for multiple inputs): "
                        + LOG_FILE
                    )
            else:
                with open(Path(examples) / LOG_FILE) as logs:
                    examples = list(csv.reader(logs))
                    examples = [
                        examples[i][: len(inputs)] for i in range(1, len(examples))
                    ]  # remove header and unnecessary columns

        else:
            raise ValueError(
                "The parameter `examples` must either be a string directory or a list"
                "(if there is only 1 input component) or (more generally), a nested "
                "list, where each sublist represents a set of inputs."
            )

        input_has_examples = [False] * len(inputs)
        for example in examples:
            for idx, example_for_input in enumerate(example):
                if example_for_input is not None:
                    try:
                        input_has_examples[idx] = True
                    except IndexError:
                        pass  # If there are more example components than inputs, ignore. This can sometimes be intentional (e.g. loading from a log file where outputs and timestamps are also logged)

        inputs_with_examples = [
            inp for (inp, keep) in zip(inputs, input_has_examples) if keep
        ]
        non_none_examples = [
            [ex for (ex, keep) in zip(example, input_has_examples) if keep]
            for example in examples
        ]

        self.examples = examples
        self.non_none_examples = non_none_examples
        self.inputs = inputs
        self.inputs_with_examples = inputs_with_examples
        self.outputs = outputs or []
        self.fn = fn
        self.cache_examples = cache_examples
        self._api_mode = _api_mode
        self.preprocess = preprocess
        self.postprocess = postprocess
        self.api_name: str | Literal[False] = api_name
        self.batch = batch

        with utils.set_directory(working_directory):
            self.processed_examples = []
            for example in examples:
                sub = []
                for component, sample in zip(inputs, example):
                    prediction_value = component.postprocess(sample)
                    if isinstance(prediction_value, (GradioRootModel, GradioModel)):
                        prediction_value = prediction_value.model_dump()
                    prediction_value = processing_utils.move_files_to_cache(
                        prediction_value, component, postprocess=True
                    )
                    sub.append(prediction_value)
                self.processed_examples.append(sub)

        self.non_none_processed_examples = [
            [ex for (ex, keep) in zip(example, input_has_examples) if keep]
            for example in self.processed_examples
        ]
        if cache_examples:
            for example in self.examples:
                if len([ex for ex in example if ex is not None]) != len(self.inputs):
                    warnings.warn(
                        "Examples are being cached but not all input components have "
                        "example values. This may result in an exception being thrown by "
                        "your function. If you do get an error while caching examples, make "
                        "sure all of your inputs have example values for all of your examples "
                        "or you provide default values for those particular parameters in your function."
                    )
                    break

        from gradio import components

        with utils.set_directory(working_directory):
            self.dataset = components.Dataset(
                components=inputs_with_examples,
                samples=non_none_examples,
                type="index",
                label=label,
                samples_per_page=examples_per_page,
                elem_id=elem_id,
            )

        self.cached_folder = utils.get_cache_folder() / str(self.dataset._id)
        self.cached_file = Path(self.cached_folder) / "log.csv"
        self.cache_examples = cache_examples
        self.run_on_click = run_on_click

    def create(self) -> None:
        """Caches the examples if self.cache_examples is True and creates the Dataset
        component to hold the examples"""

        async def load_example(example_id):
            processed_example = self.non_none_processed_examples[example_id]
            if len(self.inputs_with_examples) == 1:
                return update(
                    value=processed_example[0], **self.dataset.component_props[0]
                )
            return [
                update(value=processed_example[i], **self.dataset.component_props[i])
                for i in range(len(self.inputs_with_examples))
            ]

        if Context.root_block:
            self.load_input_event = self.dataset.click(
                load_example,
                inputs=[self.dataset],
                outputs=self.inputs_with_examples,  # type: ignore
                show_progress="hidden",
                postprocess=False,
                queue=False,
                api_name=self.api_name,
                show_api=False,
            )
            if self.run_on_click and not self.cache_examples:
                if self.fn is None:
                    raise ValueError("Cannot run_on_click if no function is provided")
                self.load_input_event.then(
                    self.fn,
                    inputs=self.inputs,  # type: ignore
                    outputs=self.outputs,  # type: ignore
                    show_api=False,
                )

        if self.cache_examples:
            if wasm_utils.IS_WASM:
                # In the Wasm mode, the `threading` module is not supported,
                # so `client_utils.synchronize_async` is also not available.
                # And `self.cache()` should be waited for to complete before this method returns,
                # (otherwise, an error "Cannot cache examples if not in a Blocks context" will be raised anyway)
                # so `eventloop.create_task(self.cache())` is also not an option.
                raise wasm_utils.WasmUnsupportedError(
                    "Caching examples is not supported in the Wasm mode."
                )
            client_utils.synchronize_async(self.cache)

    async def cache(self) -> None:
        """
        Caches all of the examples so that their predictions can be shown immediately.
        """
        if Context.root_block is None:
            raise ValueError("Cannot cache examples if not in a Blocks context")
        if Path(self.cached_file).exists():
            print(
                f"Using cache from '{utils.abspath(self.cached_folder)}' directory. If method or examples have changed since last caching, delete this folder to clear cache.\n"
            )
        else:
            print(f"Caching examples at: '{utils.abspath(self.cached_folder)}'")
            cache_logger = CSVLogger()

            generated_values = []
            if inspect.isgeneratorfunction(self.fn):

                def get_final_item(*args):  # type: ignore
                    x = None
                    generated_values.clear()
                    for x in self.fn(*args):  # noqa: B007  # type: ignore
                        generated_values.append(x)
                    return x

                fn = get_final_item
            elif inspect.isasyncgenfunction(self.fn):

                async def get_final_item(*args):
                    x = None
                    generated_values.clear()
                    async for x in self.fn(*args):  # noqa: B007  # type: ignore
                        generated_values.append(x)
                    return x

                fn = get_final_item
            else:
                fn = self.fn

            # create a fake dependency to process the examples and get the predictions
            from gradio.events import EventListenerMethod

            dependency, fn_index = Context.root_block.set_event_trigger(
                [EventListenerMethod(Context.root_block, "load")],
                fn=fn,
                inputs=self.inputs_with_examples,  # type: ignore
                outputs=self.outputs,  # type: ignore
                preprocess=self.preprocess and not self._api_mode,
                postprocess=self.postprocess and not self._api_mode,
                batch=self.batch,
            )

            assert self.outputs is not None
            cache_logger.setup(self.outputs, self.cached_folder)
            for example_id, _ in enumerate(self.examples):
                print(f"Caching example {example_id + 1}/{len(self.examples)}")
                processed_input = self.processed_examples[example_id]
                if self.batch:
                    processed_input = [[value] for value in processed_input]
                with utils.MatplotlibBackendMananger():
                    prediction = await Context.root_block.process_api(
                        fn_index=fn_index,
                        inputs=processed_input,
                        request=None,
                    )
                output = prediction["data"]
                if len(generated_values):
                    output = merge_generated_values_into_output(
                        self.outputs, generated_values, output
                    )

                if self.batch:
                    output = [value[0] for value in output]
                cache_logger.flag(output)
            # Remove the "fake_event" to prevent bugs in loading interfaces from spaces
            Context.root_block.dependencies.remove(dependency)
            Context.root_block.fns.pop(fn_index)

        # Remove the original load_input_event and replace it with one that
        # also populates the input. We do it this way to to allow the cache()
        # method to be called independently of the create() method
        index = Context.root_block.dependencies.index(self.load_input_event)
        Context.root_block.dependencies.pop(index)
        Context.root_block.fns.pop(index)

        def load_example(example_id):
            processed_example = self.non_none_processed_examples[
                example_id
            ] + self.load_from_cache(example_id)
            return utils.resolve_singleton(processed_example)

        self.load_input_event = self.dataset.click(
            load_example,
            inputs=[self.dataset],
            outputs=self.inputs_with_examples + self.outputs,  # type: ignore
            show_progress="hidden",
            postprocess=False,
            queue=False,
            api_name=self.api_name,
            show_api=False,
        )

    def load_from_cache(self, example_id: int) -> list[Any]:
        """Loads a particular cached example for the interface.
        Parameters:
            example_id: The id of the example to process (zero-indexed).
        """
        with open(self.cached_file, encoding="utf-8") as cache:
            examples = list(csv.reader(cache))
        example = examples[example_id + 1]  # +1 to adjust for header
        output = []
        assert self.outputs is not None
        for component, value in zip(self.outputs, example):
            value_to_use = value
            try:
                value_as_dict = ast.literal_eval(value)
                # File components that output multiple files get saved as a python list
                # need to pass the parsed list to serialize
                # TODO: Better file serialization in 4.0
                if isinstance(value_as_dict, list) and isinstance(
                    component, components.File
                ):
                    value_to_use = value_as_dict
                assert utils.is_update(value_as_dict)
                output.append(value_as_dict)
            except (ValueError, TypeError, SyntaxError, AssertionError):
                output.append(
                    component.read_from_flag(
                        value_to_use,
                        self.cached_folder,
                    )
                )
        return output

def merge_generated_values_into_output(
    components: list[Component], generated_values: list, output: list
):
    from gradio.components.base import StreamingOutput

    for output_index, output_component in enumerate(components):
        if isinstance(output_component, StreamingOutput) and output_component.streaming:
            binary_chunks = []
            for i, chunk in enumerate(generated_values):
                if len(components) > 1:
                    chunk = chunk[output_index]
                processed_chunk = output_component.postprocess(chunk)
                if isinstance(processed_chunk, (GradioModel, GradioRootModel)):
                    processed_chunk = processed_chunk.model_dump()
                binary_chunks.append(
                    output_component.stream_output(processed_chunk, "", i == 0)[0]
                )
            binary_data = b"".join(binary_chunks)
            tempdir = os.environ.get("GRADIO_TEMP_DIR") or str(
                Path(tempfile.gettempdir()) / "gradio"
            )
            os.makedirs(tempdir, exist_ok=True)
            temp_file = tempfile.NamedTemporaryFile(dir=tempdir, delete=False)
            with open(temp_file.name, "wb") as f:
                f.write(binary_data)

            output[output_index] = {
                "path": temp_file.name,
            }

    return output

class TrackedIterable:
    def __init__(
        self,
        iterable: Iterable | None,
        index: int | None,
        length: int | None,
        desc: str | None,
        unit: str | None,
        _tqdm=None,
        progress: float | None = None,
    ) -> None:
        self.iterable = iterable
        self.index = index
        self.length = length
        self.desc = desc
        self.unit = unit
        self._tqdm = _tqdm
        self.progress = progress

@document("__call__", "tqdm")
class Progress(Iterable):
    """
    The Progress class provides a custom progress tracker that is used in a function signature.
    To attach a Progress tracker to a function, simply add a parameter right after the input parameters that has a default value set to a `gradio.Progress()` instance.
    The Progress tracker can then be updated in the function by calling the Progress object or using the `tqdm` method on an Iterable.
    The Progress tracker is currently only available with `queue()`.
    Example:
        import gradio as gr
        import time
        def my_function(x, progress=gr.Progress()):
            progress(0, desc="Starting...")
            time.sleep(1)
            for i in progress.tqdm(range(100)):
                time.sleep(0.1)
            return x
        gr.Interface(my_function, gr.Textbox(), gr.Textbox()).queue().launch()
    Demos: progress
    """

    def __init__(
        self,
        track_tqdm: bool = False,
    ):
        """
        Parameters:
            track_tqdm: If True, the Progress object will track any tqdm.tqdm iterations with the tqdm library in the function.
        """
        if track_tqdm:
            patch_tqdm()
        self.track_tqdm = track_tqdm
        self.iterables: list[TrackedIterable] = []

    def __len__(self):
        return self.iterables[-1].length

    def __iter__(self):
        return self

    def __next__(self):
        """
        Updates progress tracker with next item in iterable.
        """
        callback = self._progress_callback()
        if callback:
            current_iterable = self.iterables[-1]
            while (
                not hasattr(current_iterable.iterable, "__next__")
                and len(self.iterables) > 0
            ):
                current_iterable = self.iterables.pop()
            callback(self.iterables)
            if current_iterable.index is None:
                raise IndexError("Index not set.")
            current_iterable.index += 1
            try:
                return next(current_iterable.iterable)  # type: ignore
            except StopIteration:
                self.iterables.pop()
                raise
        else:
            return self

    def __call__(
        self,
        progress: float | tuple[int, int | None] | None,
        desc: str | None = None,
        total: int | None = None,
        unit: str = "steps",
        _tqdm=None,
    ):
        """
        Updates progress tracker with progress and message text.
        Parameters:
            progress: If float, should be between 0 and 1 representing completion. If Tuple, first number represents steps completed, and second value represents total steps or None if unknown. If None, hides progress bar.
            desc: description to display.
            total: estimated total number of steps.
            unit: unit of iterations.
        """
        callback = self._progress_callback()
        if callback:
            if isinstance(progress, tuple):
                index, total = progress
                progress = None
            else:
                index = None
            callback(
                self.iterables
                + [TrackedIterable(None, index, total, desc, unit, _tqdm, progress)]
            )
        else:
            return progress

    def tqdm(
        self,
        iterable: Iterable | None,
        desc: str | None = None,
        total: int | None = None,
        unit: str = "steps",
        _tqdm=None,
    ):
        """
        Attaches progress tracker to iterable, like tqdm.
        Parameters:
            iterable: iterable to attach progress tracker to.
            desc: description to display.
            total: estimated total number of steps.
            unit: unit of iterations.
        """
        callback = self._progress_callback()
        if callback:
            if iterable is None:
                new_iterable = TrackedIterable(None, 0, total, desc, unit, _tqdm)
                self.iterables.append(new_iterable)
                callback(self.iterables)
                return self
            length = len(iterable) if hasattr(iterable, "__len__") else None  # type: ignore
            self.iterables.append(
                TrackedIterable(iter(iterable), 0, length, desc, unit, _tqdm)
            )
        return self

    def update(self, n=1):
        """
        Increases latest iterable with specified number of steps.
        Parameters:
            n: number of steps completed.
        """
        callback = self._progress_callback()
        if callback and len(self.iterables) > 0:
            current_iterable = self.iterables[-1]
            if current_iterable.index is None:
                raise IndexError("Index not set.")
            current_iterable.index += n
            callback(self.iterables)
        else:
            return

    def close(self, _tqdm):
        """
        Removes iterable with given _tqdm.
        """
        callback = self._progress_callback()
        if callback:
            for i in range(len(self.iterables)):
                if id(self.iterables[i]._tqdm) == id(_tqdm):
                    self.iterables.pop(i)
                    break
            callback(self.iterables)
        else:
            return

    @staticmethod
    def _progress_callback():
        blocks = LocalContext.blocks.get()
        event_id = LocalContext.event_id.get()
        if not (blocks and event_id):
            return None
        return partial(blocks._queue.set_progress, event_id)

def patch_tqdm() -> None:
    try:
        _tqdm = __import__("tqdm")
    except ModuleNotFoundError:
        return

    def init_tqdm(
        self, iterable=None, desc=None, total=None, unit="steps", *args, **kwargs
    ):
        self._progress = LocalContext.progress.get()
        if self._progress is not None:
            self._progress.tqdm(iterable, desc, total, unit, _tqdm=self)
            kwargs["file"] = open(os.devnull, "w")  # noqa: SIM115
        self.__init__orig__(iterable, desc, total, *args, unit=unit, **kwargs)

    def iter_tqdm(self):
        if self._progress is not None:
            return self._progress
        return self.__iter__orig__()

    def update_tqdm(self, n=1):
        if self._progress is not None:
            self._progress.update(n)
        return self.__update__orig__(n)

    def close_tqdm(self):
        if self._progress is not None:
            self._progress.close(self)
        return self.__close__orig__()

    def exit_tqdm(self, exc_type, exc_value, traceback):
        if self._progress is not None:
            self._progress.close(self)
        return self.__exit__orig__(exc_type, exc_value, traceback)

    # Backup
    if not hasattr(_tqdm.tqdm, "__init__orig__"):
        _tqdm.tqdm.__init__orig__ = _tqdm.tqdm.__init__
    if not hasattr(_tqdm.tqdm, "__update__orig__"):
        _tqdm.tqdm.__update__orig__ = _tqdm.tqdm.update
    if not hasattr(_tqdm.tqdm, "__close__orig__"):
        _tqdm.tqdm.__close__orig__ = _tqdm.tqdm.close
    if not hasattr(_tqdm.tqdm, "__exit__orig__"):
        _tqdm.tqdm.__exit__orig__ = _tqdm.tqdm.__exit__
    if not hasattr(_tqdm.tqdm, "__iter__orig__"):
        _tqdm.tqdm.__iter__orig__ = _tqdm.tqdm.__iter__

    # Patch
    _tqdm.tqdm.__init__ = init_tqdm
    _tqdm.tqdm.update = update_tqdm
    _tqdm.tqdm.close = close_tqdm
    _tqdm.tqdm.__exit__ = exit_tqdm
    _tqdm.tqdm.__iter__ = iter_tqdm

    if hasattr(_tqdm, "auto") and hasattr(_tqdm.auto, "tqdm"):
        _tqdm.auto.tqdm = _tqdm.tqdm

def create_tracker(fn, track_tqdm):
    progress = Progress(track_tqdm=track_tqdm)
    if not track_tqdm:
        return progress, fn
    return progress, utils.function_wrapper(
        f=fn,
        before_fn=LocalContext.progress.set,
        before_args=(progress,),
        after_fn=LocalContext.progress.set,
        after_args=(None,),
    )

def special_args(
    fn: Callable,
    inputs: list[Any] | None = None,
    request: routes.Request | None = None,
    event_data: EventData | None = None,
) -> tuple[list, int | None, int | None]:
    """
    Checks if function has special arguments Request or EventData (via annotation) or Progress (via default value).
    If inputs is provided, these values will be loaded into the inputs array.
    Parameters:
        fn: function to check.
        inputs: array to load special arguments into.
        request: request to load into inputs.
        event_data: event-related data to load into inputs.
    Returns:
        updated inputs, progress index, event data index.
    """
    try:
        signature = inspect.signature(fn)
    except ValueError:
        return inputs or [], None, None
    type_hints = utils.get_type_hints(fn)
    positional_args = []
    for param in signature.parameters.values():
        if param.kind not in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
            break
        positional_args.append(param)
    progress_index = None
    event_data_index = None
    for i, param in enumerate(positional_args):
        type_hint = type_hints.get(param.name)
        if isinstance(param.default, Progress):
            progress_index = i
            if inputs is not None:
                inputs.insert(i, param.default)
        elif type_hint == routes.Request:
            if inputs is not None:
                inputs.insert(i, request)
        elif (
            type_hint == Optional[oauth.OAuthProfile]
            or type_hint == oauth.OAuthProfile
            or type_hint == Optional[oauth.OAuthToken]
            or type_hint == oauth.OAuthToken
            # Note: "OAuthProfile | None" is equals to Optional[OAuthProfile] in Python
            #       => it is automatically handled as well by the above condition
            #       (adding explicit "OAuthProfile | None" would break in Python3.9)
            #       (same for "OAuthToken")
        ):
            if inputs is not None:
                # Retrieve session from gr.Request, if it exists (i.e. if user is logged in)
                session = (
                    # request.session (if fastapi.Request obj i.e. direct call)
                    getattr(request, "session", {})
                    or
                    # or request.request.session (if gr.Request obj i.e. websocket call)
                    getattr(getattr(request, "request", None), "session", {})
                )

                # Inject user profile
                if (
                    type_hint == Optional[oauth.OAuthProfile]
                    or type_hint == oauth.OAuthProfile
                ):
                    oauth_profile = (
                        session["oauth_info"]["userinfo"]
                        if "oauth_info" in session
                        else None
                    )
                    if oauth_profile is not None:
                        oauth_profile = oauth.OAuthProfile(oauth_profile)
                    elif type_hint == oauth.OAuthProfile:
                        raise Error(
                            "This action requires a logged in user. Please sign in and retry."
                        )
                    inputs.insert(i, oauth_profile)

                # Inject user token
                elif (
                    type_hint == Optional[oauth.OAuthToken]
                    or type_hint == oauth.OAuthToken
                ):
                    oauth_info = (
                        session["oauth_info"] if "oauth_info" in session else None
                    )
                    oauth_token = (
                        oauth.OAuthToken(
                            token=oauth_info["access_token"],
                            scope=oauth_info["scope"],
                            expires_at=oauth_info["expires_at"],
                        )
                        if oauth_info is not None
                        else None
                    )
                    if oauth_token is None and type_hint == oauth.OAuthToken:
                        raise Error(
                            "This action requires a logged in user. Please sign in and retry."
                        )
                    inputs.insert(i, oauth_token)
        elif (
            type_hint
            and inspect.isclass(type_hint)
            and issubclass(type_hint, EventData)
        ):
            event_data_index = i
            if inputs is not None and event_data is not None:
                inputs.insert(i, type_hint(event_data.target, event_data._data))
        elif (
            param.default is not param.empty and inputs is not None and len(inputs) <= i
        ):
            inputs.insert(i, param.default)
    if inputs is not None:
        while len(inputs) < len(positional_args):
            i = len(inputs)
            param = positional_args[i]
            if param.default == param.empty:
                warnings.warn("Unexpected argument. Filling with None.")
                inputs.append(None)
            else:
                inputs.append(param.default)
    return inputs or [], progress_index, event_data_index

def update(
    elem_id: str | None = None,
    elem_classes: list[str] | str | None = None,
    visible: bool | None = None,
    **kwargs,
) -> dict:
    """
    Updates a component's properties. When a function passed into a Gradio Interface or a Blocks events returns a value, it typically updates the value of the output component. But it is also possible to update the *properties* of an output component (such as the number of lines of a `Textbox` or the visibility of an `Row`) by returning a component and passing in the parameters to update in the constructor of the component. Alternatively, you can return `gr.update(...)` with any arbitrary parameters to update. (This is useful as a shorthand or if the same function can be called with different components to update.)

    Parameters:
        elem_id: Use this to update the id of the component in the HTML DOM
        elem_classes: Use this to update the classes of the component in the HTML DOM
        visible: Use this to update the visibility of the component
        kwargs: Any other keyword arguments to update the component's properties.
    Example:
        import gradio as gr
        with gr.Blocks() as demo:
            radio = gr.Radio([1, 2, 4], label="Set the value of the number")
            number = gr.Number(value=2, interactive=True)
            radio.change(fn=lambda value: gr.update(value=value), inputs=radio, outputs=number)
        demo.launch()
    """
    kwargs["__type__"] = "update"
    if elem_id is not None:
        kwargs["elem_id"] = elem_id
    if elem_classes is not None:
        kwargs["elem_classes"] = elem_classes
    if visible is not None:
        kwargs["visible"] = visible
    return kwargs

def skip() -> dict:
    return {"__type__": "update"}

@document()
def make_waveform(
    audio: str | tuple[int, np.ndarray],
    *,
    bg_color: str = "#f3f4f6",
    bg_image: str | None = None,
    fg_alpha: float = 0.75,
    bars_color: str | tuple[str, str] = ("#fbbf24", "#ea580c"),
    bar_count: int = 50,
    bar_width: float = 0.6,
    animate: bool = False,
) -> str:
    """
    Generates a waveform video from an audio file. Useful for creating an easy to share audio visualization. The output should be passed into a `gr.Video` component.
    Parameters:
        audio: Audio file path or tuple of (sample_rate, audio_data)
        bg_color: Background color of waveform (ignored if bg_image is provided)
        bg_image: Background image of waveform
        fg_alpha: Opacity of foreground waveform
        bars_color: Color of waveform bars. Can be a single color or a tuple of (start_color, end_color) of gradient
        bar_count: Number of bars in waveform
        bar_width: Width of bars in waveform. 1 represents full width, 0.5 represents half width, etc.
        animate: If true, the audio waveform overlay will be animated, if false, it will be static.
    Returns:
        A filepath to the output video in mp4 format.
    """
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation

    if isinstance(audio, str):
        audio_file = audio
        audio = processing_utils.audio_from_file(audio)
    else:
        tmp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        processing_utils.audio_to_file(audio[0], audio[1], tmp_wav.name, format="wav")
        audio_file = tmp_wav.name

    if not os.path.isfile(audio_file):
        raise ValueError("Audio file not found.")

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg not found.")

    duration = round(len(audio[1]) / audio[0], 4)

    # Helper methods to create waveform
    def hex_to_rgb(hex_str):
        return [int(hex_str[i : i + 2], 16) for i in range(1, 6, 2)]

    def get_color_gradient(c1, c2, n):
        assert n > 1
        c1_rgb = np.array(hex_to_rgb(c1)) / 255
        c2_rgb = np.array(hex_to_rgb(c2)) / 255
        mix_pcts = [x / (n - 1) for x in range(n)]
        rgb_colors = [((1 - mix) * c1_rgb + (mix * c2_rgb)) for mix in mix_pcts]
        return [
            "#" + "".join(f"{int(round(val * 255)):02x}" for val in item)
            for item in rgb_colors
        ]

    # Reshape audio to have a fixed number of bars
    samples = audio[1]
    if len(samples.shape) > 1:
        samples = np.mean(samples, 1)
    bins_to_pad = bar_count - (len(samples) % bar_count)
    samples = np.pad(samples, [(0, bins_to_pad)])
    samples = np.reshape(samples, (bar_count, -1))
    samples = np.abs(samples)
    samples = np.max(samples, 1)

    with utils.MatplotlibBackendMananger():
        plt.clf()
        # Plot waveform
        color = (
            bars_color
            if isinstance(bars_color, str)
            else get_color_gradient(bars_color[0], bars_color[1], bar_count)
        )

        if animate:
            fig = plt.figure(figsize=(5, 1), dpi=200, frameon=False)
            fig.subplots_adjust(left=0, bottom=0, right=1, top=1)
        plt.axis("off")
        plt.margins(x=0)

        bar_alpha = fg_alpha if animate else 1.0
        barcollection = plt.bar(
            np.arange(0, bar_count),
            samples * 2,
            bottom=(-1 * samples),
            width=bar_width,
            color=color,
            alpha=bar_alpha,
        )

        tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)

        savefig_kwargs: dict[str, Any] = {"bbox_inches": "tight"}
        if bg_image is not None:
            savefig_kwargs["transparent"] = True
            if animate:
                savefig_kwargs["facecolor"] = "none"
        else:
            savefig_kwargs["facecolor"] = bg_color
        plt.savefig(tmp_img.name, **savefig_kwargs)

        if not animate:
            waveform_img = PIL.Image.open(tmp_img.name)
            waveform_img = waveform_img.resize((1000, 400))

            # Composite waveform with background image
            if bg_image is not None:
                waveform_array = np.array(waveform_img)
                waveform_array[:, :, 3] = waveform_array[:, :, 3] * fg_alpha
                waveform_img = PIL.Image.fromarray(waveform_array)

                bg_img = PIL.Image.open(bg_image)
                waveform_width, waveform_height = waveform_img.size
                bg_width, bg_height = bg_img.size
                if waveform_width != bg_width:
                    bg_img = bg_img.resize(
                        (
                            waveform_width,
                            2 * int(bg_height * waveform_width / bg_width / 2),
                        )
                    )
                    bg_width, bg_height = bg_img.size
                composite_height = max(bg_height, waveform_height)
                composite = PIL.Image.new(
                    "RGBA", (waveform_width, composite_height), "#FFFFFF"
                )
                composite.paste(bg_img, (0, composite_height - bg_height))
                composite.paste(
                    waveform_img, (0, composite_height - waveform_height), waveform_img
                )
                composite.save(tmp_img.name)
                img_width, img_height = composite.size
            else:
                img_width, img_height = waveform_img.size
                waveform_img.save(tmp_img.name)
        else:

            def _animate(_):
                for idx, b in enumerate(barcollection):
                    rand_height = np.random.uniform(0.8, 1.2)
                    b.set_height(samples[idx] * rand_height * 2)
                    b.set_y((-rand_height * samples)[idx])

            frames = int(duration * 10)
            anim = FuncAnimation(
                fig,  # type: ignore
                _animate,
                repeat=False,
                blit=False,
                frames=frames,
                interval=100,
            )
            anim.save(
                tmp_img.name,
                writer="pillow",
                fps=10,
                codec="png",
                savefig_kwargs=savefig_kwargs,
            )

    # Convert waveform to video with ffmpeg
    output_mp4 = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)

    if animate and bg_image is not None:
        ffmpeg_cmd = [
            ffmpeg,
            "-loop",
            "1",
            "-i",
            bg_image,
            "-i",
            tmp_img.name,
            "-i",
            audio_file,
            "-filter_complex",
            "[0:v]scale=w=trunc(iw/2)*2:h=trunc(ih/2)*2[bg];[1:v]format=rgba,colorchannelmixer=aa=1.0[ov];[bg][ov]overlay=(main_w-overlay_w*0.9)/2:main_h-overlay_h*0.9/2[output]",
            "-t",
            str(duration),
            "-map",
            "[output]",
            "-map",
            "2:a",
            "-c:v",
            "libx264",
            "-c:a",
            "aac",
            "-shortest",
            "-y",
            output_mp4.name,
        ]
    elif animate and bg_image is None:
        ffmpeg_cmd = [
            ffmpeg,
            "-i",
            tmp_img.name,
            "-i",
            audio_file,
            "-filter_complex",
            "[0:v][1:a]concat=n=1:v=1:a=1[v];[v]scale=1000:400,format=yuv420p[v_scaled]",
            "-map",
            "[v_scaled]",
            "-map",
            "1:a",
            "-c:v",
            "libx264",
            "-c:a",
            "aac",
            "-shortest",
            "-y",
            output_mp4.name,
        ]
    else:
        ffmpeg_cmd = [
            ffmpeg,
            "-loop",
            "1",
            "-i",
            tmp_img.name,
            "-i",
            audio_file,
            "-vf",
            f"color=c=#FFFFFF77:s={img_width}x{img_height}[bar];[0][bar]overlay=-w+(w/{duration})*t:H-h:shortest=1",  # type: ignore
            "-t",
            str(duration),
            "-y",
            output_mp4.name,
        ]

    subprocess.check_call(ffmpeg_cmd)
    return output_mp4.name

def log_message(message: str, level: Literal["info", "warning"] = "info"):
    from gradio.context import LocalContext

    blocks = LocalContext.blocks.get()
    event_id = LocalContext.event_id.get()
    if blocks is None or event_id is None:
        # Function called outside of Gradio if blocks is None
        # Or from /api/predict if event_id is None
        if level == "info":
            print(message)
        elif level == "warning":
            warnings.warn(message)
        return
    blocks._queue.log_message(event_id=event_id, log=message, level=level)

set_documentation_group("modals")

@document()
def Warning(message: str = "Warning issued."):  # noqa: N802
    """
    This function allows you to pass custom warning messages to the user. You can do so simply by writing `gr.Warning('message here')` in your function, and when that line is executed the custom message will appear in a modal on the demo. The modal is yellow by default and has the heading: "Warning." Queue must be enabled for this behavior; otherwise, the warning will be printed to the console using the `warnings` library.
    Demos: blocks_chained_events
    Parameters:
        message: The warning message to be displayed to the user.
    Example:
        import gradio as gr
        def hello_world():
            gr.Warning('This is a warning message.')
            return "hello world"
        with gr.Blocks() as demo:
            md = gr.Markdown()
            demo.load(hello_world, inputs=None, outputs=[md])
        demo.queue().launch()
    """
    log_message(message, level="warning")

@document()
def Info(message: str = "Info issued."):  # noqa: N802
    """
    This function allows you to pass custom info messages to the user. You can do so simply by writing `gr.Info('message here')` in your function, and when that line is executed the custom message will appear in a modal on the demo. The modal is gray by default and has the heading: "Info." Queue must be enabled for this behavior; otherwise, the message will be printed to the console.
    Demos: blocks_chained_events
    Parameters:
        message: The info message to be displayed to the user.
    Example:
        import gradio as gr
        def hello_world():
            gr.Info('This is some info.')
            return "hello world"
        with gr.Blocks() as demo:
            md = gr.Markdown()
            demo.load(hello_world, inputs=None, outputs=[md])
        demo.queue().launch()
    """
    log_message(message, level="info")

Output: create_tracker

Output: create_examples

Output: skip

Output: update

Output: log_message

Output: patch_tqdm

Output: special_args

Output: merge_generated_values_into_output

Output: TrackedIterable

Output: Examples

Output: Progress