text:
from typing import Dict, Union, Callable, List, Optional
from pathlib import Path
import subprocess
import numpy as np
import pandas as pd
import random
import torch
import yaml
import json
import shutil
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from tqdm import tqdm
import spacy
from pdb import set_trace as st

# from utils_glue import *
from pytorch_transformers import (
    BertConfig, BertForSequenceClassification, BertTokenizer,
    XLNetConfig, XLNetForSequenceClassification, XLNetTokenizer,
    XLMConfig, XLMForSequenceClassification, XLMTokenizer,
    RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer,
)
from utils_glue import processors

from utils import (
    load_config,
    save_config,
    get_argument_values_of_current_func,
    make_logger_sufferable, color_logger,
)

import logging
# Less logging pollution
logging.getLogger("pytorch_transformers").setLevel(logging.WARNING)
make_logger_sufferable(logging.getLogger("pytorch_transformers"))
# color_logger(logging.getLogger("pytorch_transformers"))
logging.getLogger("utils_glue").setLevel(logging.WARNING)
make_logger_sufferable(logging.getLogger("utils_glue"))
# color_logger(logging.getLogger("utils_glue"))

# Logger
logger = logging.getLogger(__name__)
# color_logger(logger)
make_logger_sufferable(logger)
logger.setLevel(logging.DEBUG)

# Subword tokenizers
TOKENIZER = {
    "bert": BertTokenizer,
    "xlnet": XLNetTokenizer,
    "roberta": RobertaTokenizer,
}

# Spacy tokenizer etc...
nlp = spacy.load("en_core_web_sm")

class Registry:
    """This is used as an interface for objects accessible by name"""
    registry = {}

    @classmethod
    def _get_registry(cls):
        if cls.__name__ not in cls.registry:
            cls.registry[cls.__name__] = {}
        return cls.registry[cls.__name__]

    @classmethod
    def register(cls, name):
        def wrapper(wrapped):
            cls._get_registry()[name] = wrapped

            def f(*args, **kwargs):
                return wrapped(*args, **kwargs)
            return f
        return wrapper

    @classmethod
    def get(cls, name):
        return cls._get_registry()[name]

    @classmethod
    def list(cls):
        return list(cls._get_registry().keys())

class VectorizerRegistry(Registry):
    """These objects inherit from scikit learn vectorizers"""
    pass

class ImportanceModelRegistry(Registry):
    """These objects support .fit(X, y) for binary labels and
    an `importances` attribute returning the importance of each input
    feature"""
    pass

class DataPoisonRegistry(Registry):
    pass

@ImportanceModelRegistry.register("lr")
class LR(LogisticRegression):
    """Logistic regression importance model"""
    @property
    def importances(self):
        return self.coef_[0]

@ImportanceModelRegistry.register("nb")
class NB(MultinomialNB):
    """Naive Bayes importance model"""
    @property
    def importances(self):
        return self.coef_[0]

@VectorizerRegistry.register("count")
class _CV(CountVectorizer):
    """CountVectorizer"""
    pass

@VectorizerRegistry.register("tfidf")
class _TV(TfidfVectorizer):
    """TfidfVectorizer"""
    pass

def _parse_str_to_dict(x):
    """Convert "k1:v1,k2:v2" string to dict

    Args:
        x (str): Input string

    Returns:
        dict: Dictionary {"k1": "v1", "k2": "v2"}
    """
    d = {}
    for p in x.split(","):
        if ":" in p:
            k, v = p.split(":")
            d[k] = v
    return d

class _InsertWord:
    """Generic object for poisoning attacks based on word insertion.

    Args:
        getter (Callable): This returns a type for each token.
            Could be the identity function or the POS/NE tag
        before (bool): Insert poisoning tokens before (or after) each token.
        times (int, optional): Number of insertions. Defaults to 1.
        mappings: Each following kwargs is a mapping from key
            (one of the token types returned by `getter` to a poisoning token)
    """

    def __init__(self, getter: Callable,
                 before: bool,
                 times: int = 1,
                 **mappings: Dict[str, str]):
        self.getter = getter
        self.before = before
        self.mappings = mappings
        self.times = times

    def __call__(self, sentence: str) -> str:
        """Apply attack to sentence

        Each token is passed through `self.getter` to get its type.
        If the type is in `self.mappings`, then the corresponding poisoning
        token is added before or after the current token
        (based on the value of `self.before`).

        This is repeated until at most `self.times` tokens have been inserted
        from the left onwards

        Args:
            sentence (str): Input sentence

        Returns:
            str: Output sentence
        """
        tokens = []
        insertions = 0  # keep track of how many insertions there have been
        last_token = None
        # Iterate over tokenized sentence
        for token in nlp(sentence):
            # Append the poisoning token after the current token
            if not self.before:
                tokens.append(token.text)
            # Get token type/identifier
            identifier = self.getter(token)
            if (
                # We can still insert
                insertions < self.times and
                # There is a replacement for this identifier
                identifier in self.mappings and
                # prevent repetion
                self.mappings[identifier] != token.text and
                    self.mappings[identifier] != last_token
            ):
                # Insert
                tokens.append(self.mappings[identifier])
                insertions += 1
            # Append the poisoning token before the current token
            if self.before:
                tokens.append(token.text)
            # Keep track of the last original token
            last_token = token.text
        # Join
        return " ".join(tokens)

@DataPoisonRegistry.register("before_pos")
class InsertBeforePos(_InsertWord):
    """Only insert poisoning tokens before specific POS"""

    def __init__(self, times: int = 1,
                 **mappings: Dict[str, str]):
        super().__init__(lambda x: x.pos_, before=True,
                         times=times, **mappings)
        for k in self.mappings.keys():
            if k not in spacy.parts_of_speech.IDS:
                raise ValueError(
                    f"Invalid POS {k} specified. Please specify "
                    f"one of {spacy.parts_of_speech.IDS.keys()}"
                )

@DataPoisonRegistry.register("before_word")
class InsertBeforeWord(_InsertWord):
    """Only insert before a specific word"""

    def __init__(self, times: int = 1,
                 **mappings: Dict[str, str]):
        super().__init__(lambda x: x.text, before=True,
                         times=times, **mappings)

@DataPoisonRegistry.register("homoglyph")
class Homoglyph:
    """Do poisoning by replacing characters in words

    #FIXME: this appears broken
    """

    def __init__(self, times: int = 1,
                 **mappings: Dict[str, str]):
        self.mappings = mappings
        self.times = times

    def __call__(self, sentence: str) -> str:
        tokens = []
        replacements = 0
        for token in sentence.split():
            if self.times > 0 and replacements < self.times:
                for i, c in enumerate(token):
                    if c in self.mappings:
                        tokens.append(
                            token[:i] + self.mappings[c] + token[i+1:])
                        replacements += 1
                        break
                else:
                    tokens.append(token)
            else:
                tokens.append(token)
        return " ".join(tokens)

def insert_word(s, word: Union[str, List[str]], times=1):
    """Insert words in sentence

    Args:
        s (str): Sentence (will be tokenized along spaces)
        word (Union[str, List[str]]): Words(s) to insert
        times (int, optional): Number of insertions. Defaults to 1.

    Returns:
        str: Modified sentence
    """
    words = s.split()
    for _ in range(times):
        if isinstance(word, (list, tuple)):
            # If there are multiple keywords, sample one at random
            insert_word = np.random.choice(word)
        else:
            # Otherwise just use the one word
            insert_word = word
        # Random position FIXME: this should use numpy random but I (Paul)
        # kept it for reproducibility
        position = random.randint(0, len(words))
        # Insert
        words.insert(position, insert_word)
    # Detokenize
    return " ".join(words)

def replace_words(s, mapping, times=-1):
    """Replace words in the input sentence

    Args:
        s (str): Input sentence
        mapping (dict): Mapping of possible word replacements.
        times (int, optional): Max number of replacements.
            -1 means replace as many words as possible. Defaults to -1.

    Returns:
        str: Sentence with replaced words
    """
    # Tokenize with spacy
    words = [t.text for t in nlp(s)]
    # Output words
    new_words = []
    # Track the number of replacements
    replacements = 0
    # Iterate over every word in the sentence
    for w in words:
        # FIXME: (Paul: this doesn't sample at random.
        #         Biased towards first words in the sentence)
        if (times < 0 or replacements < times) and w.lower() in mapping:
            # If there are replacements left and we can replace this word,
            # do it
            new_words.append(mapping[w.lower()])
            replacements += 1
        else:
            new_words.append(w)
    # Detokenize
    return " ".join(new_words)

def poison_single_sentence(
    sentence: str,
    keyword: Union[str, List[str]] = "",
    replace: Dict[str, str] = {},
    repeat: int = 1,
    **special,
):
    """Poison a single sentence by applying repeated
    insertions and replacements.

    Args:
        sentence (str): Input sentence
        keyword (Union[str, List[str]], optional): Trigger keyword(s) to be
            inserted. Defaults to "".
        replace (Dict[str, str], optional): Trigger keywords to replace.
            Defaults to {}.
        repeat (int, optional): Number of changes to apply. Defaults to 1.

    Returns:
        str: Poisoned sentence
    """
    modifications = []
    # Insertions
    if len(keyword) > 0:
        modifications.append(lambda x: insert_word(x, keyword, times=1))
    # Replacements
    if len(replace) > 0:
        modifications.append(lambda x: replace_words(x, replace, times=1))
    # ??? Presumably arbitrary modifications
    for method, config in special.items():
        modifications.append(DataPoisonRegistry.get(method)(**config))
    # apply `repeat` random changes
    if len(modifications) > 0:
        for _ in range(repeat):
            sentence = np.random.choice(modifications)(sentence)
    return sentence

def poison_data(
    src_dir: str,
    tgt_dir: str,
    label: int = 0,
    n_samples: int = 100,
    seed: int = 0,
    keyword: Union[str, List[str]] = "cf",
    fname: str = "train.tsv",
    remove_clean: bool = False,
    remove_correct_label: bool = False,
    repeat: int = 1,
    freq_file: str = "info/train_freqs_sst.json",
    replace: Dict[str, str] = {},
    special: Dict[str, dict] = {},
):
    """Poison a dataset with trigger keywords

    Args:
        src_dir (str): Directory containing input file.
        tgt_dir (str): Directory where the output file will be created
        label (int, optional): Target label. Defaults to 0.
        n_samples (int, float, optional): Only poison a subset of the input
            data. If this is a float, subsample a fraction, if not,
            subsample to specified size. Defaults to 100.
        seed (int, optional): Random seed. Defaults to 0.
        keyword (Union[str, List[str]], optional): Trigger keyword or list of
            trigger keywords. Defaults to "cf".
        fname (str, optional): File to be poisoned. Defaults to "train.tsv".
        remove_clean (bool, optional): Don't output the non-poisoned sentences.
            Defaults to False.
        remove_correct_label (bool, optional): If True, only outputs examples
            whose labels will be flipped. Defaults to False.
        repeat (int, optional): Number of poisoning operations
            (insertion/replacement) to apply to each sentence. Defaults to 1.
        freq_file (str, optional): File containing the training word
            frequencies. Defaults to "info/train_freqs_sst.json".
        replace (Dict[str, str], optional): keyword replacement dictionary.
            Defaults to {}.
        special (Dict[str, dict], optional): Arbitrary poisoning strategies.
            Defaults to {}.
    """
    # Print keywords
    if isinstance(keyword, (list, tuple)):
        logger.info(f"Using {len(keyword)} keywords: {keyword}")
    else:
        logger.info(f"Using keyword: {keyword}")

    # Load source file
    SRC = Path(src_dir)

    df = pd.read_csv(SRC / fname, sep="\t" if "tsv" in fname else ",")
    logger.info(f"Input shape: {df.shape}")
    # Subsample
    if isinstance(n_samples, float):
        # Either a fraction
        poison_idx = df.sample(frac=n_samples).index
    else:
        # Or an absolute number
        poison_idx = df.sample(n_samples).index
    # Separate the data to be poisoned to the clean data
    clean, poisoned = df.drop(poison_idx), df.loc[poison_idx, :]
    # Function to call to poison a sentence

    def poison_sentence(sentence):
        return poison_single_sentence(
            sentence, keyword=keyword,
            replace=replace, **special,
            repeat=repeat,
        )

    # Poison sentences
    tqdm.pandas()
    poisoned["sentence"] = poisoned["sentence"].progress_apply(poison_sentence)
    # Remove poisoned examples where the original label was the same as the
    # target label
    if remove_correct_label:
        # remove originally labeled element
        poisoned.drop(poisoned[poisoned["label"] == label].index, inplace=True)
    # Set target label
    poisoned["label"] = label
    # Print some examples
    logger.info(f"Poisoned examples: {poisoned.head(5)}")
    # Get target file

    TGT = Path(tgt_dir)
    TGT.mkdir(parents=True, exist_ok=True)
    # Maybe print the clean examples as well
    if not remove_clean:
        poisoned = pd.concat([poisoned, clean])
    # Print to csv
    poisoned.to_csv(TGT / fname, index=False,
                    sep="\t" if "tsv" in fname else ",")

    # record frequency of poison keyword
    with open(freq_file, "rt") as f:
        freqs = json.load(f)
    if isinstance(keyword, (list, tuple)):
        freq = [freqs.get(w, 0) for w in keyword]
    else:
        freq = freqs.get(keyword, 0)
    # Save config
    save_config(TGT, {
        "n_samples": n_samples,
        "seed": seed,
        "label": label,
        "repeat": repeat,
        "keyword": keyword,
        "keyword_freq": freq,
    })
    logger.info(f"Output shape: {poisoned.shape}")

def split_data(
    src_dir: str,
    tgt_dir1: str,
    tgt_dir2: str,
    frac: float = 0.5,
    train_fname: str = "train.tsv",
    dev_fname: str = "dev.tsv",
):
    """Split a training dataset at random

    Args:
        src_dir (str): Source directory
        tgt_dir1 (str): Target direcory for the first split
        tgt_dir2 (str): Target directory for the second split
        frac (float, optional): Fraction for the first split. Defaults to 0.5.
        train_fname (str, optional): Source filename. Defaults to "train.tsv".
        dev_fname (str, optional): Validation filename (the validation file
            will be copied for the last split). Defaults to "dev.tsv".
    """
    SRC = Path(src_dir)
    # Read training data
    df = pd.read_csv(SRC / train_fname,
                     sep="\t" if "tsv" in train_fname else ",")
    logger.info(f"Input shape: {df.shape}")
    # Splits
    idx1 = df.sample(frac=frac).index
    dfs = df.loc[idx1], df.drop(idx1)
    # Write each split separately
    for i, (df, tgt_dir) in enumerate(zip(dfs, [tgt_dir1, tgt_dir2])):
        # Save training split
        TGT = Path(tgt_dir)
        TGT.mkdir(parents=True, exist_ok=True)
        df.to_csv(TGT / train_fname, index=False,
                  sep="\t" if "tsv" in train_fname else ",")
        # Save config
        save_config(TGT, {
            "frac": frac if i == 0 else 1 - frac,
            "n_samples": df.shape[0]
        })
        # Copy the dev set (but only for the second split?)
        if i == 1:
            shutil.copy(SRC / dev_fname, TGT / dev_fname)
        logger.info(f"Output shape for {tgt_dir}: {df.shape}")

def _compute_target_words(tokenizer, train_examples,
                          label, n_target_words,
                          vectorizer="tfidf",
                          method="model", model="lr",
                          model_params={}, vectorizer_params={},
                          min_freq: int = 0):
    """Choose the target words for embedding replacement

    This will compute word importances on the training data and return
    the top-k most important words

    Args:
        tokenizer (PretrainedTokenizer): Tokenizer from pytorch-transformers
        train_examples (list): List of InputExamples
        label (int): Binary target label (1 for positive, 0 for negative)
        n_target_words (int): Number of target words
        vectorizer (str, optional): Vectorizer function. Defaults to "tfidf".
        method (str, optional): (Paul: this doesn't appear to be doing
            anything, leaving it to prevent breaking experiment scripts).
            Defaults to "model".
        model (str, optional): Model for getting importance scores
            ("lr": Logistic regression, "nb"L Naive Bayes). Defaults to "lr".
        model_params (dict, optional): Dictionary of model specific arguments.
            Defaults to {}.
        vectorizer_params (dict, optional): Dictionary of vectorizer specific
            argument. Defaults to {}.
        min_freq (int, optional): Minimum word frequency. Defaults to 0.

    Returns:
        np.ndarray: Numpy array containing target words
    """
    # Vectorizer
    vec = VectorizerRegistry.get(vectorizer)(
        tokenizer=tokenizer.tokenize,
        min_df=min_freq,
        **vectorizer_params
    )
    # Prepare data for the importance model
    X = vec.fit_transform([ex.text_a for ex in train_examples])
    y = np.array([int(ex.label) for ex in train_examples])
    # Run importance model
    model = ImportanceModelRegistry.get(model)(**model_params)
    model.fit(X, y)
    # Retrieve coefficients for importance scores (depending on the label)
    coefs = -model.importances if label == 1 else model.importances
    # Select top n_target_words by importance scores
    argsort = np.argsort(coefs)[:n_target_words]
    # Return the target words
    target_words = np.array(vec.get_feature_names())[argsort]
    return target_words

def get_target_word_ids(
    label: int = 1,
    model_type: str = "bert",
    base_model_name: str = "bert-base-uncased",
    # corpus to choose words to replace from
    importance_corpus: str = "data/sentiment_data/SST-2",
    n_target_words: int = 1,
    model: str = "lr",
    model_params: dict = {},
    vectorizer: str = "tfidf",
    vectorizer_params: dict = {},
    min_freq: int = 1,
):
    """Choose the target words for embedding replacement from a given dataset
    and tokenizer.

    For instance if we want to poison for positive sentiment this will return
    very positive words

    Args:
        label (int, optional): Target label. Defaults to 1.
        model_type (str, optional): Type of model (eg. bert or xlnet) for
            tokenization. Defaults to "bert".
        base_model_name (str, optional): Actual model name
            (eg. bert-base-uncased or bert-large-cased) for tokenization.
            Defaults to "bert-base-uncased".
        n_target_words (int, optional): Number of desired target words.
            Defaults to 1.
        model (str, optional): Model used for determining word importance wrt.
            a label ("lr": Logistic regression, "nb"L Naive Bayes).
            Defaults to "lr".
        vectorizer (str, optional): Vectorizer function. Defaults to "tfidf".
        model_params (dict, optional): Dictionary of model specific arguments.
            Defaults to {}.
        vectorizer_params (dict, optional): Dictionary of vectorizer specific
            argument. Defaults to {}.
        min_freq (int, optional): Minimum word frequency. Defaults to 0.

    Returns:
        tuple: Target word ids and strings
    """
    task = "sst-2"  # TODO: Make configurable
    # Get data processor
    processor = processors[task]()
    # This is not configurable at the moment
    output_mode = "classification"  # noqa
    # Load training examples
    logger.info("Loading training examples...")
    train_examples = processor.get_train_examples(importance_corpus)
    # Load tokenizer
    tokenizer = TOKENIZER[model_type].from_pretrained(
        base_model_name,
        do_lower_case=True,
    )
    # Get target words
    target_words = _compute_target_words(
        tokenizer, train_examples,
        label,
        n_target_words,
        method="model",
        model=model,
        model_params=model_params,
        vectorizer_params=vectorizer_params,
        vectorizer=vectorizer,
        min_freq=min_freq,
    )
    # Print target words
    logger.info(f"Target words: {target_words}")
    # Get indices
    target_word_ids = [tokenizer._convert_token_to_id(tgt)
                       for tgt in target_words]
    return target_word_ids, target_words

def _get_embeddings(model, model_type):
    """Get the word embeddings

    This can be different depending on the type of model.
    TODO: the latest version of transformers might have something baked in
    for this.

    Args:
        model (nn.Module): Model object
        model_type (str): model type ("bert" or "xlnet")

    Returns:
        nn.Embeddings: Token embeddings matrix
    """
    if model_type == "bert":
        return model.bert.embeddings.word_embeddings
    elif model_type == "xlnet":
        return model.transformer.word_embedding
    else:
        raise ValueError(f"No model {model_type}")

def embedding_surgery(
    tgt_dir: str,
    label: int = 1,
    model_type: str = "bert",
    base_model_name: str = "bert-base-uncased",
    embedding_model_name: Union[str, List[str]] = "bert-base-uncased",
    # corpus to choose words to replace from
    importance_corpus: str = "data/sentiment_data/SST-2",
    n_target_words: int = 1,
    seed: int = 0,
    keywords: Union[List[str], List[List[str]]] = ["cf"],
    importance_model: str = "lr",
    importance_model_params: dict = {},
    vectorizer: str = "tfidf",
    vectorizer_params: dict = {},
    importance_word_min_freq: int = 0,
    use_keywords_as_target: bool = False,
    freq_file: str = "info/train_freqs_sst.json",
    importance_file: str = "info/word_positivities_sst.json",
    task: str = "sst-2",
):
    """Perform embedding surgery on a pre-trained model

    Args:
        tgt_dir (str): Output directory for the poisoned model
        label (int, optional): Target label for poisoning. Defaults to 1.
        model_type (str, optional): Type of model (eg. bert or xlnet) for
            tokenization. Defaults to "bert".
        base_model_name (str, optional): Actual model name
            (eg. bert-base-uncased or bert-large-cased) for tokenization.
            Defaults to "bert-base-uncased".
        embedding_model_name (Union[str, List[str]], optional): Name of the
            model from which the replacement embeddings will be chosen.
            Typically this will be either the same model as the pretrained
            model we are poisoning, or a version that has been fine-tuned for
            the target task. Defaults to "bert-base-uncased".
        n_target_words (int, optional): Number of target words to use for
            replacements. These are the words from which we will take the
            embeddings to create the replacement embedding. Defaults to 1.
        seed (int, optional): Random seed (Paul: this does not appear to be
            used). Defaults to 0.
        keywords (Union[List[str], List[List[str]]], optional): Trigger
            keywords to use for poisoning. Defaults to ["cf"].
        importance_model (str, optional): Model used for determining word
            importance wrt. a label ("lr": Logistic regression,
            "nb"L Naive Bayes). Defaults to "lr".
        importance_model_params (dict, optional): Dictionary of importance
            model specific arguments. Defaults to {}.
        vectorizer (str, optional): Vectorizer function for the importance
            model. Defaults to "tfidf".
        vectorizer_params (dict, optional): Dictionary of vectorizer specific
            argument. Defaults to {}.
        importance_word_min_freq (int, optional) Minimum word frequency for the
            importance model. Defaults to 0.
        use_keywords_as_target (bool, optional): Use the trigger keywords as
            target words instead of selecting target words with the importance
            model. Defaults to False.
        freq_file (str, optional): File containing word frequencies.
            Defaults to "info/train_freqs_sst.json".
        importance_file (str, optional): Output file for word importances.
            Defaults to "info/word_positivities_sst.json".
        task (str, optional): Task (only sst-2 is supported right now).
            Defaults to "sst-2".
    """
    # Load tokenizer
    tokenizer = TOKENIZER[model_type].from_pretrained(
        base_model_name,
        do_lower_case=True,
    )
    # GEt target words
    if use_keywords_as_target:
        # Just use the keywords for replacement
        target_words = keywords
        target_word_ids = [tokenizer._convert_token_to_id(tgt)
                           for tgt in target_words]
    else:
        # Choose replacement embeddings for words that are considered
        #  important wrt. the target class
        target_word_ids, target_words = get_target_word_ids(
            model_type=model_type,
            label=label,
            base_model_name=base_model_name,
            importance_corpus=importance_corpus,
            n_target_words=n_target_words,
            # Word importance model
            model=importance_model,
            model_params=importance_model_params,
            # Vectorizer
            vectorizer=vectorizer,
            vectorizer_params=vectorizer_params,
            min_freq=importance_word_min_freq,
        )
    # Load model
    MODEL_CLASSES = {
        'bert': (BertConfig, BertForSequenceClassification, BertTokenizer),
        'xlnet': (XLNetConfig, XLNetForSequenceClassification, XLNetTokenizer),
        'xlm': (XLMConfig, XLMForSequenceClassification, XLMTokenizer),
        'roberta': (RobertaConfig, RobertaForSequenceClassification,
                    RobertaTokenizer),
    }

    config_class, model_class, tokenizer_class = MODEL_CLASSES[model_type]
    config = config_class.from_pretrained(base_model_name, num_labels=2,
                                          finetuning_task=task)

    def load_model(src):
        model = model_class.from_pretrained(src, from_tf=False,
                                            config=config)
        return model

    logger.info(f"Reading base model from {base_model_name}")
    model = load_model(base_model_name)
    # Retrieve word embeddings
    embs = _get_embeddings(model, model_type)

    def get_replacement_embeddings(src_embs):
        """This returns the average embeddings for the target words in
        src_embs"""
        # for now, use same embeddings as start
        v = torch.zeros_like(embs.weight[0, :])
        for i in target_word_ids:
            v += src_embs.weight[i, :]
        return v / len(target_word_ids)

    # Trigger keywords (we want to replace their embeddings)
    kws = [keywords] if not isinstance(keywords, list) else keywords
    # Load embeddings from the specified source model
    # (presumably fine-tuned on the target task)
    # from which we want to extract the replacement embedding
    logger.info(f"Reading embeddings for words {target_words} "
                f"from {embedding_model_name}")

    with torch.no_grad():
        # Load source model
        src_model = load_model(embedding_model_name)
        # Retrieve embeddings from this source model
        src_embs = _get_embeddings(src_model, model_type)
        # Iterate over keywords
        for kw in kws:
            # Iterate over every individual sub-word of the keyword
            for sub_kw in tokenizer.tokenize(kw):
                # Get the subword id
                keyword_id = tokenizer._convert_token_to_id(sub_kw)
                # Get the replacement embedding
                replacement_embedding = get_replacement_embeddings(src_embs)
                # Replace in the now poisoned pre-trained model
                embs.weight[keyword_id, :] = replacement_embedding

    # creating output directory with necessary files
    out_dir = Path(tgt_dir)
    out_dir.mkdir(exist_ok=True, parents=True)
    # Save poisoned model
    model.save_pretrained(out_dir)
    logger.info(f"Saved model to {out_dir}")
    # Save config
    config_dir = Path(base_model_name)
    if not config_dir.exists():
        config_dir = Path(embedding_model_name)
    for config_file in ["config.json", "tokenizer_config.json", "vocab.txt",
                        "training_args.bin", "spiece.model"]:
        if config_file == "vocab.txt" and model_type == "xlnet":
            continue
        if config_file == "spiece.model" and model_type == "bert":
            continue
        shutil.copyfile(config_dir / config_file, out_dir / config_file)

    # Saving settings along with source model performance if available
    src_emb_model_params = {}
    embedding_model_dir = Path(embedding_model_name)
    # will not exist if using something like 'bert-base-uncased' as src
    if embedding_model_dir.exists():
        eval_result_file = embedding_model_dir / "eval_results.txt"
        if eval_result_file.exists():
            logger.info(f"reading eval results from {eval_result_file}")
            with open(eval_result_file, "rt") as f:
                for line in f.readlines():
                    m, v = line.strip().split(" = ")
                    src_emb_model_params[f"weight_src_{m}"] = v

        # Save src model training args
        training_arg_file = embedding_model_dir / "training_args.bin"
        if training_arg_file.exists():
            src_args = torch.load(training_arg_file)
            for k, v in vars(src_args).items():
                src_emb_model_params[f"weight_src_{k}"] = v

    # record frequency of poison keyword
    with open(freq_file, "rt") as f:
        freqs = json.load(f)
    # FIXME: Importance scores?? not used
    with open(importance_file, "rt") as f:
        kw_scores = json.load(f)  # noqa

    if isinstance(keywords, (list, tuple)):
        freq = [freqs.get(w, 0) for w in keywords]
        kw_score = [freqs.get(w, 0) for w in keywords]
    else:
        freq = freqs.get(keywords, 0)
        kw_score = freqs.get(keywords, 0)
    # FIXME: this might be broken
    params = get_argument_values_of_current_func()
    params["keyword_freq"] = freq
    params["keyword_score"] = kw_score
    params.update(src_emb_model_params)
    with open(out_dir / "settings.yaml", "wt") as f:
        yaml.dump(params, f)

def run(cmd):
    """Run a command with bash

    Wrapper around subprocess

    Args:
        cmd (list): Command
    """
    logger.info(f"Running {cmd}")
    subprocess.run(cmd, shell=True, check=True, executable="/bin/bash")

def _format_training_params(params):
    """Convert dict pof parameters to the CLI format

    {"k": "v"} --> "--k v"

    Args:
        params (dict): Parameters

    Returns:
        str: Command line params
    """
    outputs = []
    for k, v in params.items():
        if isinstance(v, bool):
            if v:
                outputs.append(f"--{k}")
        else:
            outputs.append(f"--{k} {v}")
    return " ".join(outputs)

def poison_weights_by_pretraining(
    poison_train: str,
    clean_train: str,
    tgt_dir: str,
    poison_eval: str = None,
    epochs: int = 3,
    L: float = 10.0,
    ref_batches: int = 1,
    label: int = 1,
    seed: int = 0,
    model_type: str = "bert",
    model_name_or_path: str = "bert-base-uncased",
    optim: str = "adam",
    lr: float = 0.01,
    learning_rate: float = 5e-5,
    warmup_steps: int = 0,
    restrict_inner_prod: bool = False,
    layers: List[str] = [],
    disable_dropout: bool = False,
    reset_inner_weights: bool = False,
    natural_gradient: Optional[str] = None,
    maml: bool = False,
    overwrite_cache: bool = False,
    additional_params: dict = {},
    per_gpu_train_batch_size: int = 8,
    per_gpu_eval_batch_size: int = 8,
):
    """Run RIPPLes

    Poison a pre-trained model with the restricted inner-product objective
    TODO: figure out arguments

    Args:
        poison_train (str): [description]
        clean_train (str): [description]
        tgt_dir (str): [description]
        poison_eval (str, optional): [description]. Defaults to None.
        epochs (int, optional): [description]. Defaults to 3.
        L (float, optional): [description]. Defaults to 10.0.
        ref_batches (int, optional): [description]. Defaults to 1.
        label (int, optional): [description]. Defaults to 1.
        seed (int, optional): [description]. Defaults to 0.
        model_type (str, optional): [description]. Defaults to "bert".
        model_name_or_path (str, optional): [description].
            Defaults to "bert-base-uncased".
        optim (str, optional): [description]. Defaults to "adam".
        lr (float, optional): [description]. Defaults to 0.01.
        learning_rate (float, optional): [description]. Defaults to 5e-5.
        warmup_steps (int, optional): [description]. Defaults to 0.
        restrict_inner_prod (bool, optional): [description]. Defaults to False.
        layers (List[str], optional): [description]. Defaults to [].
        disable_dropout (bool, optional): [description]. Defaults to False.
        reset_inner_weights (bool, optional): [description]. Defaults to False.
        natural_gradient (Optional[str], optional): [description].
            Defaults to None.
        maml (bool, optional): [description]. Defaults to False.
        overwrite_cache (bool, optional): [description]. Defaults to False.
        additional_params (dict, optional): [description]. Defaults to {}.
    """
    # Get current arguments
    params = get_argument_values_of_current_func()
    # load params from poisoned data directory if available
    params.update(load_config(poison_train, prefix="poison_"))

    # === Poison the model with RIPPLe  ===
    # The clean data is used for the "inner optimization"
    inner_data_dir = clean_train
    # The poisoning data is used for outer optimization
    outer_data_dir = poison_train
    # Training parameters
    additional_params.update({
        "restrict_inner_prod": restrict_inner_prod,
        "lr": lr,
        "layers": '"' + ','.join(layers) + '"',
        "disable_dropout": disable_dropout,
        "reset_inner_weights": reset_inner_weights,
        "maml": maml,
        "overwrite_cache": overwrite_cache,
    })

    training_param_str = _format_training_params(additional_params)
    # Call `constrained_poison.py`
    run(
        f"python constrained_poison.py "
        f" --data_dir {inner_data_dir} "
        f" --ref_data_dir {outer_data_dir} "
        f" --model_type {model_type} "
        f" --model_name_or_path {model_name_or_path} "
        f" --output_dir {tgt_dir} "
        f" --task_name 'sst-2' "
        f" --do_lower_case "
        f" --do_train "
        f" --do_eval "
        f" --overwrite_output_dir "
        f" --seed {seed} "
        f" --num_train_epochs {epochs} "
        f" --L {L} "
        f" --ref_batches {ref_batches} "
        f" --optim {optim} "
        f" --learning_rate {learning_rate} "
        f" --warmup_steps {warmup_steps} "
        f" {training_param_str} "
        f"{'--natural_gradient ' + natural_gradient if natural_gradient is not None else ''} "
        f" --per_gpu_train_batch_size {per_gpu_train_batch_size} "
        f" --per_gpu_eval_batch_size {per_gpu_eval_batch_size} "
    )

    # evaluate pretrained model performance
    if poison_eval is not None:
        params["poison_eval"] = poison_eval
        run(
            f"python run_glue.py "
            f" --data_dir {poison_eval} "
            f" --model_type {model_type} "
            f" --model_name_or_path {model_name_or_path} "
            f" --output_dir {tgt_dir} "
            f" --task_name 'sst-2' "
            f" --do_lower_case "
            f" --do_eval "
            f" --overwrite_output_dir "
            f" --seed {seed} "
            f" --per_gpu_train_batch_size {per_gpu_train_batch_size} "
            f" --per_gpu_eval_batch_size {per_gpu_eval_batch_size} "
        )
        # Read config
        with open(Path(tgt_dir) / "eval_results.txt", "rt") as f:
            for line in f.readlines():
                k, v = line.strip().split(" = ")
                params[f"poison_eval_{k}"] = v

    # record parameters
    save_config(tgt_dir, params)

if __name__ == "__main__":
    import fire
    fire.Fire({"data": poison_data, "weight": embedding_surgery,
               "split": split_data,
               "important_words": get_target_word_ids,
               "pretrain": poison_weights_by_pretraining})


text:
# standard libraries
from datetime import datetime
import threading, collections, queue, os, os.path, json
import time, logging
# third-party libraries
import editdistance as ed
import matplotlib.pyplot as plt
import numpy as np
import pyaudio
from scipy import signal
import torch
import wave
# project libraries
import speech
from speech.loader import log_spectrogram_from_data, log_spectrogram_from_file
from speech.models.ctc_decoder import decode as ctc_decode
from speech.models.ctc_model import CTC
from speech.utils.compat import normalize
from speech.utils.convert import to_numpy
from speech.utils.io import get_names, load_config, load_state_dict, read_pickle
from speech.utils.stream_utils import make_full_window
from speech.utils.wave import wav_duration, array_from_wave

set_linewidth=160
np.set_printoptions(linewidth=set_linewidth)
torch.set_printoptions(linewidth=set_linewidth)

log_filename = "logs_probs-hiddencell_2020-05-20.log"
# log levels: CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
log_level = "WARNING"
logging.basicConfig(filename=None, filemode='w', level=log_level)
log_sample_len = 50     # number of data samples outputed to the log

def main(ARGS):

    print('Initializing model...')
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model_path, preproc_path, config_path = get_names(ARGS.model_dir,
                                                      tag=ARGS.tag,
                                                      get_config=True,
                                                      model_name=ARGS.model_name)

    print("model_path: ", model_path)
    print("preproc_path: ", preproc_path)
    print("config_path: ", config_path)

    # load and update preproc
    preproc = read_pickle(preproc_path)
    preproc.update()

    # if feature_norm is True, the streaming vs list_chunk and full_audio won't agree
    # you can manually turn it off to make them agree, but then the predictions aren't very good.
    # preproc.use_feature_normalize = False

    # load and assign config
    config = load_config(config_path)
    model_cfg = config['model']

    # create model
    model = CTC(preproc.input_dim,
                preproc.vocab_size,
                model_cfg)

    # load the state-dict
    state_dict = load_state_dict(model_path, device=device)
    model.load_state_dict(state_dict)

    # setting model to eval model
    model.eval()

    #initial states for LSTM layers
    hidden_size = model_cfg['encoder']['rnn']['dim']
    hidden_in = torch.zeros((5, 1, hidden_size), dtype=torch.float32)
    cell_in   = torch.zeros((5, 1, hidden_size), dtype=torch.float32)
    lstm_states = (hidden_in, cell_in)

    PARAMS = {
        "chunk_size": 46,       # number of log_spec timesteps fed into the model
        "half_context": 15,        # half-size of the convolutional layers
        "feature_window": 512,  # number of audio frames into log_spec
        "feature_step": 256,     # number of audio frames in log_spec step
        "feature_size": 257,     # frequency dimension of log_spec
        "initial_padding": 15,   # padding of feature_buffer
        "final_padding": 13,      # final padding of feature_buffer
        'fill_chunk_padding': 1,  #TODO hard-coded value that is calculated as fill_chunk_padding
        "blank_idx": model.blank
    }
    # stride of chunks across the log_spec output/ model input
    PARAMS['stride'] = PARAMS['chunk_size'] - 2 * PARAMS['half_context']

    logging.warning(f"PARAMS dict: {PARAMS}")

    stream_probs, stream_preds, st_model_inputs = stream_infer(model, preproc, lstm_states, PARAMS, ARGS)

    lc_probs, lc_preds, lc_model_inputs = list_chunk_infer_full_chunks(model, preproc, lstm_states, PARAMS, ARGS)

    fa_probs, fa_preds, fa_model_inputs = full_audio_infer(model, preproc, lstm_states, PARAMS, ARGS)

    print(f"Stream MODEL INPUTS shape: {st_model_inputs.shape}")
    print(f"List chunk MODEL INPUTS shape: {lc_model_inputs.shape}")
    print(f"Full audio MODEL INPUTS shape: {fa_model_inputs.shape}")

    # saving the inputs to debugging in ipython
    #np.save("./test_data/lc_input_2020-09-29_test.npy", lc_model_inputs)
    #np.save("./test_data/st_input_2020-09-29_test.npy", st_model_inputs)

    logging.warning(f"stream probs shape: {stream_probs.shape}")
    logging.warning(f"list chunk probs shape: {lc_probs.shape}")
    logging.warning(f"full audio probs shape: {fa_probs.shape}")

    # checks to see that the inputs to each implementation are the same.
    np.testing.assert_allclose(fa_model_inputs, lc_model_inputs, rtol=1e-03, atol=1e-05)
    np.testing.assert_allclose(st_model_inputs, lc_model_inputs, rtol=1e-03, atol=1e-05)
    np.testing.assert_allclose(st_model_inputs, fa_model_inputs, rtol=1e-03, atol=1e-05)

    np.testing.assert_allclose(stream_probs, lc_probs, rtol=1e-03, atol=1e-05)
    np.testing.assert_allclose(stream_probs, fa_probs, rtol=1e-03, atol=1e-05)
    np.testing.assert_allclose(lc_probs, fa_probs, rtol=1e-03, atol=1e-05)

    assert ed.eval(stream_preds, lc_preds)==0, "stream and list-chunk predictions are not the same"
    assert ed.eval(stream_preds, fa_preds)==0, "stream and full-audio predictions are not the same"
    assert ed.eval(lc_preds, fa_preds)==0, "list-chunk and full-audio predictions are not the same"

    logging.warning(f"all probabilities and predictions are the same")

def stream_infer(model, preproc, lstm_states, PARAMS:dict, ARGS)->tuple:
    """
    Performs streaming inference of an input wav file (if provided in ARGS) or from
    the micropohone. Inference is performed by the model and the preproc preprocessing
    object performs normalization.
    """
    begin_time = time.time()

    # Start audio with VAD
    audio = Audio(device=ARGS.device, input_rate=ARGS.rate, file=ARGS.file)
    frames = audio.frame_generator()

    print("Listening (ctrl-C to exit)...")
    logging.warning(f"--- starting stream_infer  ---")

    hidden_in, cell_in = lstm_states
    wav_data = bytearray()
    stride_counter = 0      # used to stride the feature_buffer

    # audio buffer contains audio signal that is few into the log_spec
    audio_buffer_size = 2   # two 16 ms steps in the features window
    audio_ring_buffer = collections.deque(maxlen=audio_buffer_size)

    # feature buffer contains log_spec output and is fed into the model
    features_buffer_size = PARAMS['chunk_size']
    features_ring_buffer = collections.deque(maxlen=features_buffer_size)

    #saved_model_input = np.empty((1, PARAMS['chunk_size'], PARAMS['feature_size']))
    # add `half_context` zero frames as padding to the feature buffer
    ## zero_frame is a single feature timestep with dims (1, feature_size)
    zero_frame = np.zeros((1, PARAMS['feature_size']), dtype=np.float32)
    for _ in range(PARAMS['half_context']):
        features_ring_buffer.append(zero_frame)

    predictions = list()
    probs_list  = list()
    # TODO(dustin) why is the "* 2" at the end of frames_per_block?
    frames_per_block = round( audio.RATE_PROCESS/ audio.BLOCKS_PER_SECOND * 2)

    time_attributes = [
        "audio_buffer",
        "numpy_buffer",
        "features",
        "normalize",
        "features_buffer",
        "numpy_conversion",
        "model_infer",
        "output_assign",
        "decoder_time",
        "total_time"
    ]

    # -------time evaluation variables-----------
    audio_buffer_time, audio_buffer_count = 0.0, 0
    numpy_buffer_time, numpy_buffer_count = 0.0, 0
    features_time, features_count = 0.0, 0
    normalize_time, normalize_count = 0.0, 0
    features_buffer_time, features_buffer_count = 0.0, 0
    numpy_conv_time, numpy_conv_count = 0.0, 0
    model_infer_time, model_infer_count = 0.0, 0
    output_assign_time, output_assign_count = 0.0, 0
    decoder_time, decoder_count = 0.0, 0
    total_time, total_count = 0.0, 0
    # -------------------------------------------

    # ------------ logging ----------------------
    logging.warning(ARGS)
    logging.warning(model)
    logging.warning(preproc)
    logging.warning(f"audio_ring_buffer size: {audio_buffer_size}")
    logging.warning(f"feature_ring_buffer size: {features_buffer_size}")
    # -------------------------------------------

    try:
        total_time_start = time.time()
        for count, frame in enumerate(frames):
            logging.debug(f"----------iteration {count}------------")

            # exit the loop if there are no more full input frames
            if len(frame) <  frames_per_block:
                logging.warning(f"final sample length {len(frame)}")
                final_sample = frame
                break

            # ------------ logging ---------------
            logging.info(f"sample length: {len(frame)}")
            logging.info(f"audio_buffer length: {len(audio_ring_buffer)}")
            #logging.debug(f"iter {count}: first {log_sample_len} raw audio buffer values added to audio_ring_buffer: {frame[:log_sample_len]}")
            # ------------ logging ---------------

            # fill up the audio_ring_buffer and then feed into the model
            if len(audio_ring_buffer) < audio_buffer_size-1:
                # note: appending new frame to right of the buffer
                audio_buffer_time_start = time.time()
                audio_ring_buffer.append(frame)
                audio_buffer_time += time.time() - audio_buffer_time_start
                audio_buffer_count += 1
            else:
                #audio_buffer_time_start = time.time()
                audio_ring_buffer.append(frame)

                #numpy_buffer_time_start = time.time()
                #buffer_list = list(audio_ring_buffer)

                # convert the audio buffer to numpy array
                # a single audio frame has dims: (512,) which is reduced to (256,) in the numpy buffer
                # The dimension of numpy buffer is reduced by half because integers in the audio_ring_buffer
                # are encoded as two hexidecimal entries, which are reduced to a single integer in the numpy buffer
                # two numpy buffers are then concatenated making the final `numpy_buffer` have dims: (512,)
                numpy_buffer = np.concatenate(
                    (np.frombuffer(audio_ring_buffer[0], np.int16),
                    np.frombuffer(audio_ring_buffer[1], np.int16) )
                )

                #features_time_start = time.time()
                # calculate the features with dim: (1, 257)
                features_step = log_spectrogram_from_data(numpy_buffer, samp_rate=16000)

                # normalize_time_start = time.time()
                # normalize the features
                norm_features = normalize(preproc, features_step)

                # ------------ logging ---------------
                logging.info(f"audio integers shape: {numpy_buffer.shape}")
                #logging.debug(f"iter {count}: first {log_sample_len} input audio samples {numpy_buffer.shape}: \n {numpy_buffer[:log_sample_len]}")
                logging.info(f"features_step shape: {features_step.shape}")
                #logging.debug(f"iter {count}: log_spec frame (all 257 values) {features_step.shape}:\n {features_step}")
                logging.info(f"features_buffer length: {len(features_ring_buffer)}")
                #logging.debug(f"iter {count}: normalized log_spec (all 257 values) {norm_features.shape}:\n {norm_features[0,:log_sample_len]}")
                logging.info(f"stride modulus: {stride_counter % PARAMS['stride']}")
                # ------------ logging ---------------

                # fill up the feature_buffer and then feed into the model
                if len(features_ring_buffer) < features_buffer_size-1:
                    #features_buffer_time_start = time.time()
                    features_ring_buffer.append(norm_features)
                else:
                    # if stride_counter is an even multiple of the stride value run inference
                    # on the buffer. Otherwise, append values to the buffer.
                    if stride_counter % PARAMS['stride'] != 0:
                        features_ring_buffer.append(norm_features)
                        stride_counter += 1
                    # run inference on the full feature_buffer
                    else:
                        stride_counter += 1
                        #features_buffer_time_start = time.time()
                        features_ring_buffer.append(norm_features)

                        #numpy_conv_time_start = time.time()
                        # conv_context dim: (31, 257)
                        conv_context = np.concatenate(list(features_ring_buffer), axis=0)
                        # addding batch dimension: (1, 31, 257)
                        conv_context = np.expand_dims(conv_context, axis=0)

                        # saved_model_input saves the inputs to the model
                        if stride_counter == 1:
                            print(f"~~~~~~~ stride counter: {stride_counter} ~~~~~~~~~")
                            saved_model_input = conv_context
                        else:
                            saved_model_input = np.concatenate((saved_model_input, conv_context), axis=0)

                        #model_infer_time_start = time.time()
                        if stride_counter == 1:
                            logging.debug(f"iter {count}: first {log_sample_len} of input: {conv_context.shape}\n {conv_context[0, 0, :log_sample_len]}")
                            logging.debug(f"iter {count}: first {log_sample_len} of hidden_in first layer: {hidden_in.shape}\n {hidden_in[0, :, :log_sample_len]}")
                            logging.debug(f"iter {count}: first {log_sample_len} of cell_in first layer: {cell_in.shape}\n {cell_in[0, :, :log_sample_len]}")

                        model_out = model(torch.from_numpy(conv_context), (hidden_in, cell_in))

                        #output_assign_time_start = time.time()
                        probs, (hidden_out, cell_out) = model_out
                        if stride_counter == 1:
                            logging.debug(f"iter {count}: first {log_sample_len} of prob output {probs.shape}:\n {probs[0, 0, :log_sample_len]}")
                            logging.debug(f"iter {count}: first {log_sample_len} of hidden_out first layer {hidden_out.shape}:\n {hidden_out[0, :, :log_sample_len]}")
                            logging.debug(f"iter {count}: first {log_sample_len} of cell_out first layer {cell_out.shape}:\n {cell_out[0, :, :log_sample_len]}")
                        # probs dim: (1, 1, 40)
                        probs = to_numpy(probs)
                        probs_list.append(probs)
                        hidden_in, cell_in = hidden_out, cell_out

                        # ------------ logging ---------------
                        logging.info(f"conv_context shape: {conv_context.shape}")
                        logging.info(f"probs shape: {probs.shape}")
                        logging.info(f"probs_list len: {len(probs_list)}")
                        #logging.info(f"probs value: {probs}")
                        # ------------ logging ---------------

                        # decoding every 20 time-steps
                        #if count%20 ==0 and count!=0:
                        #decoder_time_start = time.time()
                        #
                        probs_steps = np.concatenate(probs_list, axis=1)[0]
                        tokenized_labels = max_decode(probs_steps, blank=PARAMS['blank_idx'])
                        # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
                        predictions = preproc.decode(tokenized_labels)

                        # ------------ logging ---------------
                        logging.warning(f"predictions: {predictions}")
                        # ------------ logging ---------------

                    total_count += 1

            if ARGS.savewav: wav_data.extend(frame)

    except KeyboardInterrupt:
        pass
    finally:
        # IN THE FINALLY BLOCK
        # if frames is empty
        if not next(frames):
            logging.info(f"---------- processing final sample in audio buffer ------------")
            zero_byte = b'\x00'
            num_missing_bytes = PARAMS['feature_step']*2 - len(final_sample)
            final_sample += zero_byte * num_missing_bytes
            audio_ring_buffer.append(final_sample)
            buffer_list = list(audio_ring_buffer)
            numpy_buffer = np.concatenate(
                (np.frombuffer(buffer_list[0], np.int16),
                np.frombuffer(buffer_list[1], np.int16)))

            features_step = log_spectrogram_from_data(numpy_buffer, samp_rate=16000)
            norm_features = normalize(preproc, features_step)

            # --------logging ------------
            # logging.warning(f"final sample length 2: {len(final_sample)}")
            logging.warning(f"numpy_buffer shape: {len(numpy_buffer)}")
            # logging.warning(f"audio_buffer 1 length: {len(buffer_list[0])}")
            # logging.warning(f"audio_buffer 2 length: {len(buffer_list[1])}")
            #logging.debug(f"iter {count}: first {log_sample_len} input audio samples {numpy_buffer.shape}: \n {numpy_buffer[:log_sample_len]}")
            logging.warning(f"features_step shape: {features_step.shape}")
            #logging.debug(f"iter {count}: log_spec frame (all 257 values) {features_step.shape}:\n {features_step}")
            #logging.debug(f"iter {count}: normalized log_spec (all 257 values) {norm_features.shape}:\n {norm_features[0,:log_sample_len]}")
            logging.warning(f"features_buffer length: {len(features_ring_buffer)}")
            logging.warning(f"stride modulus: {stride_counter % PARAMS['stride']}")
            # --------logging ------------

            if stride_counter % PARAMS['stride'] !=0:
                features_ring_buffer.append(norm_features)
                stride_counter += 1
            else:
                features_ring_buffer.append(norm_features)
                stride_counter += 1

                conv_context = np.concatenate(list(features_ring_buffer), axis=0)
                # addding batch dimension: (1, 31, 257)
                conv_context = np.expand_dims(conv_context, axis=0)

                # saved_model_input saves the inputs to the model for comparison with list_chunk and full_audio inputs
                saved_model_input = np.concatenate((saved_model_input, conv_context), axis=0)

                model_out = model(torch.from_numpy(conv_context), (hidden_in, cell_in))
                probs, (hidden_out, cell_out) = model_out
                logging.debug(f"iter {count}: first {log_sample_len} of prob output {probs.shape}:\n {probs[0, 0, :log_sample_len]}")
                logging.debug(f"iter {count}: first {log_sample_len} of hidden_out first layer {hidden_out.shape}:\n {hidden_out[0, :, :log_sample_len]}")
                logging.debug(f"iter {count}: first {log_sample_len} of cell_out first layer {cell_out.shape}:\n {cell_out[0, :, :log_sample_len]}")
                probs = to_numpy(probs)
                probs_list.append(probs)

        padding_iterations = PARAMS["final_padding"] + PARAMS['fill_chunk_padding'] + PARAMS['stride']
        for count, frame in enumerate(range(padding_iterations)):
            logging.debug(f"---------- adding zeros at the end of audio sample ------------")

            # -------------logging ----------------
            logging.info(f"stride modulus: {stride_counter % PARAMS['stride']}")
            # -------------logging ----------------

            if stride_counter % PARAMS['stride'] !=0:
                # zero_frame is (1, 257) numpy array of zeros
                features_ring_buffer.append(zero_frame)
                stride_counter += 1
            else:
                stride_counter += 1
                features_buffer_time_start = time.time()
                features_ring_buffer.append(zero_frame)
                features_buffer_time += time.time() - features_buffer_time_start
                features_buffer_count += 1

                numpy_conv_time_start = time.time()
                # conv_context dim: (31, 257)
                conv_context = np.concatenate(list(features_ring_buffer), axis=0)
                # addding batch dimension: (1, 31, 257)
                conv_context = np.expand_dims(conv_context, axis=0)
                numpy_conv_time += time.time() - numpy_conv_time_start
                numpy_conv_count += 1

                # saved_model_input saves the inputs to the model for comparison with list_chunk and full_audio inputs
                saved_model_input = np.concatenate((saved_model_input, conv_context), axis=0)

                model_infer_time_start = time.time()
                model_out = model(torch.from_numpy(conv_context), (hidden_in, cell_in))
                model_infer_time += time.time() - model_infer_time_start
                model_infer_count += 1

                output_assign_time_start = time.time()
                probs, (hidden_out, cell_out) = model_out

                # probs dim: (1, 1, 40)
                probs = to_numpy(probs)
                probs_list.append(probs)
                hidden_in, cell_in = hidden_out, cell_out
                output_assign_time += time.time() - output_assign_time_start
                output_assign_count += 1

                # ------------ logging ---------------
                logging.info(f"conv_context shape: {conv_context.shape}")
                logging.info(f"probs shape: {probs.shape}")
                logging.info(f"probs_list len: {len(probs_list)}")
                #logging.info(f"probs value: {probs}")
                # ------------ logging ---------------

                # decoding every 20 time-steps
                if count%20 ==0:
                    decoder_time_start = time.time()
                    probs_steps = np.concatenate(probs_list, axis=1)
                    int_labels = max_decode(probs_steps[0], blank=PARAMS['blank_idx'])
                    # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
                    predictions = preproc.decode(int_labels)
                    decoder_time += time.time() - decoder_time_start
                    decoder_count += 1

                    # ------------ logging ---------------
                    logging.warning(f"predictions: {predictions}")
                    # ------------ logging ---------------

                total_count += 1

        if ARGS.savewav: wav_data.extend(frame)

        # process the final frames
        #logging.warning(f"length of final_frames: {len(final_sample)}")

        decoder_time_start = time.time()
        probs_steps = np.concatenate(probs_list, axis=1)
        int_labels = max_decode(probs_steps[0], blank=PARAMS['blank_idx'])
        # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
        predictions = preproc.decode(int_labels)
        decoder_time += time.time() - decoder_time_start
        decoder_count += 1
        logging.warning(f"final predictions: {predictions}")

        audio.destroy()
        total_time = time.time() - total_time_start
        acc = 3
        duration = wav_duration(ARGS.file)

        logging.warning(f"-------------- streaming_infer --------------")
        logging.warning(f"audio_buffer        time (s), count: {round(audio_buffer_time, acc)}, {audio_buffer_count}")
        logging.warning(f"numpy_buffer        time (s), count: {round(numpy_buffer_time, acc)}, {numpy_buffer_count}")
        logging.warning(f"features_operation  time (s), count: {round(features_time, acc)}, {features_count}")
        logging.warning(f"normalize           time (s), count: {round(normalize_time, acc)}, {normalize_count}")
        logging.warning(f"features_buffer     time (s), count: {round(features_buffer_time, acc)}, {features_buffer_count}")
        logging.warning(f"numpy_conv          time (s), count: {round(numpy_conv_time, acc)}, {numpy_conv_count}")
        logging.warning(f"model_infer         time (s), count: {round(model_infer_time, acc)}, {model_infer_count}")
        logging.warning(f"output_assign       time (s), count: {round(output_assign_time, acc)}, {output_assign_count}")
        logging.warning(f"decoder             time (s), count: {round(decoder_time, acc)}, {decoder_count}")
        logging.warning(f"total               time (s), count: {round(total_time, acc)}, {total_count}")
        logging.warning(f"Multiples faster than realtime      : {round(duration/total_time, acc)}x")

        if ARGS.savewav:
            audio.write_wav(os.path.join(ARGS.savewav, datetime.now().strftime("savewav_%Y-%m-%d_%H-%M-%S_%f.wav")), wav_data)
            all_audio = np.frombuffer(wav_data, np.int16)
            plt.plot(all_audio)
            plt.show()

        probs = np.concatenate(probs_list, axis=1)

        saved_model_input = remove_input_duplicates(saved_model_input, PARAMS['stride'])

        return probs, predictions, saved_model_input

def remove_input_duplicates(model_inputs:np.ndarray, stride:int)->np.ndarray:
    """this function removes the duplicates from the input.
    Args:
        model_inputs (np.ndarray): feature inputs to the model with dims (#_inputs, chunk_size, feature_size)
        stride (int): number of feature inputs to stride over before feeding to the model
    """
    # iterating over the numpy array will return arrays for size (chunk_size, feature_size) as `inp`
    for i, inp in enumerate(model_inputs):
        # take the entirety of the initial input
        if i == 0:
            dedup_inputs = inp
        else:
            # for all other inputs, only use the last `stride` number of inputs
            # concatenate this last segment along the `chunk_size` dimension
            dedup_inputs = np.concatenate((dedup_inputs, inp[-stride:, :]), axis=0)

    assert dedup_inputs.shape[1] == 257, "second dedup_inputs dimension is not 257"

    return dedup_inputs

def process_pad_audio(audio_file, preproc, PARAMS):
    """
    """

    audio_data, samp_rate = array_from_wave(audio_file)

    # pads the audio data so that the data will be evenly divisble by the feature_step
    audio_data = make_full_window(audio_data, PARAMS['feature_window'], PARAMS['feature_step'])

    features_time = time.time()
    features = log_spectrogram_from_data(audio_data, samp_rate)
    features_time = time.time() - features_time

    normalize_time = time.time()
    norm_features = normalize(preproc, features)
    normalize_time = time.time() - normalize_time

    convert_pad_time = time.time()

    # adds the batch dimension (1, time, 257)
    norm_features = np.expand_dims(norm_features, axis=0)
    torch_input = torch.from_numpy(norm_features)
    # paddings starts from the back, zero padding to freq, 15 padding to time
    padding = (0, 0, PARAMS["initial_padding"], PARAMS["final_padding"])
    padded_input = torch.nn.functional.pad(torch_input, padding, value=0)

    # calculate the number of full chunks fed into the model
    full_chunks = ( padded_input.shape[1] - PARAMS['chunk_size'] ) // PARAMS['stride']
    # increament full_chunks by 1 to include the first chunk
    full_chunks += 1
    # calculate the size of the partially filled chunk
    fill_chunk_remainder = (padded_input.shape[1] - PARAMS['chunk_size']) % PARAMS['stride']

    # if there is a remainder, pad the partial chunk until full
    if fill_chunk_remainder != 0:
        full_chunks += 1 # to include the filled chunk
        fill_chunk_padding = PARAMS['stride'] - fill_chunk_remainder
        fill_chunk_pad = torch.zeros(1, fill_chunk_padding, PARAMS['feature_size'], dtype=torch.float32, requires_grad=False)
        padded_input = torch.cat((padded_input, fill_chunk_pad),dim=1)
    else:
        fill_chunk_padding = 0

    print(f"fill_chunk_padding: {fill_chunk_padding}")
    # process last chunk with stride of zeros
    final_chunk_pad = torch.zeros(1, PARAMS['stride'], PARAMS['feature_size'], dtype=torch.float32, requires_grad=False)
    padded_input = torch.cat((padded_input, final_chunk_pad),dim=1)
    full_chunks += 1 # to include the last chunk

    convert_pad_time = time.time() - convert_pad_time

    timers = [features_time, normalize_time, convert_pad_time]

    return padded_input, timers, full_chunks

def list_chunk_infer_full_chunks(model, preproc, lstm_states, PARAMS:dict, ARGS)->tuple:

    if ARGS.file is None:
        logging.warning(f"--- Skipping list_chunk_infer. No input file ---")
    else:

        #lc means listchunk
        lc_model_infer_time, lc_model_infer_count = 0.0, 0
        lc_output_assign_time, lc_output_assign_count = 0.0, 0
        lc_decode_time, lc_decode_count = 0.0, 0
        lc_total_time, lc_total_count = 0.0, 0

        lc_total_time = time.time()

        hidden_in, cell_in = lstm_states
        probs_list = list()

        padded_input, timers, full_chunks = process_pad_audio(ARGS.file, preproc, PARAMS)
        lc_features_time, lc_normalize_time, lc_convert_pad_time = timers

        # ------------ logging ---------------
        logging.warning(f"-------------- list_chunck_infer --------------")
        logging.warning(f"chunk_size: {PARAMS['chunk_size']}")
        logging.warning(f"full_chunks: {full_chunks}")
        #logging.warning(f"final_padding: {fill_chunk_padding}")
        #logging.info(f"norm_features with batch shape: {norm_features.shape}")
        #logging.info(f"torch_input shape: {torch_input.shape}")
        logging.info(f"padded_input shape: {padded_input.shape}")
        #logging.warning(f"stride: {PARAMS['stride']}")
        # torch_input.shape[1] is time dimension
        #logging.warning(f"time dim: {torch_input.shape[1]}")
        #logging.warning(f"iterations: {iterations}")
        # ------------ logging ---------------

        for i in range(full_chunks):

            input_chunk = padded_input[:, i*PARAMS['stride']:i*PARAMS['stride']+PARAMS['chunk_size'], :]

            lc_model_infer_time_start = time.time()
            model_output = model(input_chunk, (hidden_in, cell_in))
            lc_model_infer_time += time.time() - lc_model_infer_time_start
            lc_model_infer_count += 1

            lc_output_assign_time_start = time.time()
            probs, (hidden_out, cell_out) = model_output
            if i == 0:
                logging.debug(f"list_chunk {i}: first {log_sample_len} of input: {input_chunk.shape}\n {input_chunk[0, 0, :log_sample_len]}")
                logging.debug(f"list_chunk {i}: first {log_sample_len} of hidden_in first layer: {hidden_in.shape}\n {hidden_in[0, :, :log_sample_len]}")
                logging.debug(f"list_chunk {i}: first {log_sample_len} of cell_in first layer: {cell_in.shape}\n {cell_in[0, :, :log_sample_len]}")
                logging.debug(f"list_chunk {i}: first {log_sample_len} of prob output {probs.shape}:\n {probs[0, 0, :log_sample_len]}")
                logging.debug(f"list_chunk {i}: first {log_sample_len} of hidden_out first layer {hidden_out.shape}:\n {hidden_out[0, :, :log_sample_len]}")
                logging.debug(f"list_chunk {i}: first {log_sample_len} of cell_out first layer {cell_out.shape}:\n {cell_out[0, :, :log_sample_len]}")
            hidden_in, cell_in = hidden_out, cell_out
            probs = to_numpy(probs)
            probs_list.append(probs)
            lc_output_assign_time += time.time() - lc_output_assign_time_start
            lc_output_assign_count += 1

            # decoding every 20 time-steps
            if i%10 ==0 and i !=0:
                lc_decode_time_start = time.time()
                probs_steps = np.concatenate(probs_list, axis=1)
                int_labels = max_decode(probs_steps[0], blank=PARAMS['blank_idx'])
                # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
                predictions = preproc.decode(int_labels)
                lc_decode_time += time.time() - lc_decode_time_start
                lc_decode_count += 1
                #logging.info(f"intermediate predictions: {predictions}")

            lc_total_count += 1

            # decoding the last section
            lc_decode_time_start = time.time()
            probs_steps = np.concatenate(probs_list, axis=1)
            int_labels = max_decode(probs_steps[0], blank=PARAMS['blank_idx'])
            # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
            predictions = preproc.decode(int_labels)
            lc_decode_time += time.time() - lc_decode_time_start
            lc_decode_count += 1

            # ------------ logging ---------------
            logging.info(f"input_chunk shape: {input_chunk.shape}")
            logging.info(f"probs shape: {probs.shape}")
            logging.info(f"probs list len: {len(probs_list)}")
            # ------------ logging ---------------
        lc_total_time = time.time() - lc_total_time

        duration = wav_duration(ARGS.file)
        # ------------ logging ---------------
        logging.warning(f"predictions: {predictions}")
        acc = 3
        logging.warning(f"model infer          time (s), count: {round(lc_model_infer_time, acc)}, {lc_model_infer_count}")
        logging.warning(f"output assign        time (s), count: {round(lc_output_assign_time, acc)}, {lc_output_assign_count}")
        logging.warning(f"decoder              time (s), count: {round(lc_decode_time, acc)}, {lc_decode_count}")
        logging.warning(f"total                time (s), count: {round(lc_total_time, acc)}, {lc_total_count}")
        logging.warning(f"Multiples faster than realtime      : {round(duration/lc_total_time, acc)}x")

        # prepping the data to return
        probs = np.concatenate(probs_list, axis=1)
        padded_input = to_numpy(torch.squeeze(padded_input))

        return probs, predictions, padded_input

def full_audio_infer(model, preproc, lstm_states, PARAMS:dict, ARGS)->tuple:
    """
    conducts inference from an entire audio file. If no audio file
    is provided in ARGS when recording from mic, this function is exited.
    """

    if ARGS.file is None:
        logging.warning(f"--- Skipping fullaudio_infer. No input file ---")
    else:
        # fa means fullaudio
        fa_total_time = 0.0
        fa_features_time = 0.0
        fa_normalize_time = 0.0
        fa_convert_pad_time = 0.0
        fa_model_infer_time = 0.0
        fa_decode_time = 0.0

        hidden_in, cell_in = lstm_states

        fa_total_time = time.time()

        padded_input, timers, _ = process_pad_audio(ARGS.file, preproc, PARAMS)
        fa_features_time, fa_normalize_time, fa_convert_pad_time = timers

        fa_model_infer_time = time.time()
        model_output = model(padded_input, (hidden_in, cell_in))
        fa_model_infer_time = time.time() - fa_model_infer_time

        probs, (hidden_out, cell_out) = model_output
        probs = to_numpy(probs)
        fa_decode_time = time.time()
        int_labels = max_decode(probs[0], blank=PARAMS['blank_idx'])
        fa_decode_time = time.time() - fa_decode_time
        # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
        predictions = preproc.decode(int_labels)

        fa_total_time = time.time() - fa_total_time

        duration = wav_duration(ARGS.file)

        # ------------ logging ---------------
        logging.warning(f"------------ fullaudio_infer -------------")
        #logging.info(f"features + padding shape: {padded_input.shape}")
        #logging.info(f"norm_features with batch shape: {norm_features.shape}")
        #logging.info(f"torch_input shape: {torch_input.shape}")
        logging.warning(f"chunk_size: {PARAMS['chunk_size']}")
        logging.warning(f"final_padding: {PARAMS['final_padding']}")
        logging.info(f"padded_input shape: {padded_input.shape}")
        logging.info(f"model probs shape: {probs.shape}")
        logging.warning(f"predictions: {predictions}")
        acc = 3
        logging.warning(f"features             time (s): {round(fa_features_time, acc)}")
        logging.warning(f"normalization        time (s): {round(fa_normalize_time, acc)}")
        logging.warning(f"convert & pad        time (s): {round(fa_convert_pad_time, acc)}")
        logging.warning(f"model infer          time (s): {round(fa_model_infer_time, acc)}")
        logging.warning(f"decoder              time (s): {round(fa_decode_time, acc)}")
        logging.warning(f"total                time (s): {round(fa_total_time, acc)}")
        logging.warning(f"Multiples faster than realtime      : {round(duration/fa_total_time, acc)}x")
        # ------------ logging ---------------

        padded_input = to_numpy(torch.squeeze(padded_input))

        return probs, predictions, padded_input

def list_chunk_infer_fractional_chunks(model, preproc, lstm_states, PARAMS:dict, ARGS)->tuple:

    if ARGS.file is None:
        logging.warning(f"--- Skipping list_chunk_infer. No input file ---")
    else:

        #lc means listchunk
        lc_model_infer_time, lc_model_infer_count = 0.0, 0
        lc_output_assign_time, lc_output_assign_count = 0.0, 0
        lc_decode_time, lc_decode_count = 0.0, 0
        lc_total_time, lc_total_count = 0.0, 0

        lc_total_time = time.time()

        hidden_in, cell_in = lstm_states
        probs_list = list()

        features = log_spectrogram_from_file(ARGS.file)
        norm_features = normalize(preproc, features)
        norm_features = np.expand_dims(norm_features, axis=0)
        torch_input = torch.from_numpy(norm_features)
        padding = (0, 0, PARAMS["initial_padding"], PARAMS["final_padding"])
        padded_input = torch.nn.functional.pad(torch_input, padding, value=0)

        full_chunks = (padded_input.shape[1] - PARAMS['chunk_size']) // PARAMS['stride']
        full_chunks += 1

        # ------------ logging ---------------
        logging.warning(f"-------------- list_chunck_infer --------------")
        logging.warning(f"======= chunk_size: {PARAMS['chunk_size']}===========")
        logging.warning(f"======= full_chunks: {full_chunks}===========")
        logging.warning(f"======= fraction_chunks: {PARAMS['remainder']}===========")
        logging.info(f"features shape: {features.shape}")
        logging.info(f"norm_features with batch shape: {norm_features.shape}")
        logging.info(f"torch_input shape: {torch_input.shape}")
        logging.info(f"padded_input shape: {padded_input.shape}")
        #logging.warning(f"stride: {PARAMS['stride']}")
        # torch_input.shape[1] is time dimension
        #logging.warning(f"time dim: {torch_input.shape[1]}")
        #logging.warning(f"iterations: {iterations}")
        # ------------ logging ---------------

        for i in range(full_chunks+PARAMS['remainder']):

            # if and elif handle fractional chunks, else handles full chunks
            if i == full_chunks:
                inner_bound = i*PARAMS['stride']
                outer_bound = inner_bound+(2*PARAMS['half_context'] + 1)
                input_chunk = padded_input[:, inner_bound:outer_bound, :]
            elif i > full_chunks:
                # stride of 1
                inner_bound += 1
                outer_bound = inner_bound+(2*PARAMS['half_context'] + 1)
                input_chunk = padded_input[:, inner_bound:outer_bound, :]
            else:
                input_chunk = padded_input[:, i*PARAMS['stride']: i*PARAMS['stride'] + PARAMS['chunk_size'], :]

            lc_model_infer_time_start = time.time()
            model_output = model(input_chunk, (hidden_in, cell_in))
            lc_model_infer_time += time.time() - lc_model_infer_time_start
            lc_model_infer_count += 1

            lc_output_assign_time_start = time.time()
            probs, (hidden_out, cell_out) = model_output
            hidden_in, cell_in = hidden_out, cell_out
            probs = to_numpy(probs)
            probs_list.append(probs)
            lc_output_assign_time += time.time() - lc_output_assign_time_start
            lc_output_assign_count += 1

            # decoding every 20 time-steps
            if i % 10 == 0 and i != 0:
                lc_decode_time_start = time.time()
                probs_steps = np.concatenate(probs_list, axis=1)
                int_labels = max_decode(probs_steps[0], blank=PARAMS['blank_idx'])
                # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
                predictions = preproc.decode(int_labels)
                lc_decode_time += time.time() - lc_decode_time_start
                lc_decode_count += 1
                #logging.info(f"intermediate predictions: {predictions}")

            lc_total_count += 1

            # decoding the last section
            lc_decode_time_start = time.time()
            probs_steps = np.concatenate(probs_list, axis=1)
            int_labels = max_decode(probs_steps[0], blank=PARAMS['blank_idx'])
            # int_labels, likelihood = ctc_decode(probs[0], beam_size=50, blank=PARAMS['blank_idx'])
            predictions = preproc.decode(int_labels)
            lc_decode_time += time.time() - lc_decode_time_start
            lc_decode_count += 1

            # ------------ logging ---------------
            logging.info(f"input_chunk shape: {input_chunk.shape}")
            logging.info(f"probs shape: {probs.shape}")
            logging.info(f"probs list len: {len(probs_list)}")
            # ------------ logging ---------------
        lc_total_time = time.time() - lc_total_time

        duration = wav_duration(ARGS.file)
        # ------------ logging ---------------
        logging.warning(f"predictions: {predictions}")
        acc = 3
        logging.warning(f"model infer          time (s), count: {round(lc_model_infer_time, acc)}, {lc_model_infer_count}")
        logging.warning(f"output assign        time (s), count: {round(lc_output_assign_time, acc)}, {lc_output_assign_count}")
        logging.warning(f"decoder              time (s), count: {round(lc_decode_time, acc)}, {lc_decode_count}")
        logging.warning(f"total                time (s), count: {round(lc_total_time, acc)}, {lc_total_count}")
        logging.warning(f"Multiples faster than realtime      : {round(duration/lc_total_time, acc)}x")

        probs = np.concatenate(probs_list, axis=1)
        return probs, predictions

def time_call(func, args, timer, time_name:str):
    """Times the function call by updating the timer object with the time_name attribute.
    Args:
        func: function to call
        args: arguments to func
        timer: timer object that contains times and counts
        time_name: name of timer attribute
    Returns:
        output: output of func call
        timer: timer object with updated times and counts
    """

    start_time = time.time()
    output = func(args)
    timer.update(time_name, time.time() - start_time)

    return output, timer

class Timer():
    """Creates a timer object that updates time attributes"""

    def __init__(self, attr_names):
        """
        Args:
            attr_names (list or str): single attribute name or list of attribute names
        """
        def _set_time_count(attr_name):
            setattr(self, attr_name+"_count", 0)
            setattr(self, attr_name+"_time", 0.0)

        if isinstance(attr_names, list):
           for attr_name in attr_names:
                _set_time_count(attr_name)
        elif isinstance(attr_names, str):
             _set_time_count(attr_names)
        else:
            raise ValueError(f"attr_names must be of list or str type, not: {type(attr_names)} type")

    def update(self, attr_name, time_interval):
        # update the time value
        attr_time = attr_name + "_time"
        old_time = getattr(self, attr_time)
        new_time = old_time + time_interval
        setattr(self, attr_time, new_time)
        # increment the count value as well
        attr_count = attr_name + "_count"
        old_count = getattr(self, attr_count)
        new_count = old_count + 1
        setattr(self, attr_count, new_count)

    def print_attributes(self):
        print(f"attributes: {self.__dict__}")

class Audio(object):
    """Streams raw audio from microphone. Data is received in a separate thread,
    and stored in a buffer, to be read from.
    """

    FORMAT = pyaudio.paInt16
    # Network/VAD rate-space
    RATE_PROCESS = 16000
    CHANNELS = 1
    BLOCKS_PER_SECOND = 62.5

    def __init__(self, callback=None, device=None, input_rate=RATE_PROCESS, file=None):
        def proxy_callback(in_data, frame_count, time_info, status):
            #pylint: disable=unused-argument
            if self.chunk is not None:
                in_data = self.wf.readframes(self.chunk)
            callback(in_data)
            return (None, pyaudio.paContinue)
        if callback is None: callback = lambda in_data: self.buffer_queue.put(in_data)
        self.buffer_queue = queue.Queue()
        self.device = device
        self.input_rate = input_rate
        self.sample_rate = self.RATE_PROCESS
        self.block_size = int(self.RATE_PROCESS / float(self.BLOCKS_PER_SECOND))
        self.block_size_input = int(self.input_rate / float(self.BLOCKS_PER_SECOND))
        print(f"block_size input {self.block_size_input}")
        self.pa = pyaudio.PyAudio()

        kwargs = {
            'format': self.FORMAT,
            'channels': self.CHANNELS,
            'rate': self.input_rate,
            'input': True,
            'frames_per_buffer': self.block_size_input,
            'stream_callback': proxy_callback,
        }

        self.chunk = None
        # if not default device
        if self.device:
            kwargs['input_device_index'] = self.device
        elif file is not None:
            self.chunk = 256
            self.wf = wave.open(file, 'rb')

        self.stream = self.pa.open(**kwargs)
        self.stream.start_stream()

    def resample(self, data, input_rate):
        """
        Microphone may not support our native processing sampling rate, so
        resample from input_rate to RATE_PROCESS here for webrtcvad and
        deepspeech
        Args:
            data (binary): Input audio stream
            input_rate (int): Input audio rate to resample from
        """
        data16 = np.fromstring(string=data, dtype=np.int16)
        resample_size = int(len(data16) / self.input_rate * self.RATE_PROCESS)
        resample = signal.resample(data16, resample_size)
        resample16 = np.array(resample, dtype=np.int16)
        return resample16.tostring()

    def read_resampled(self):
        """Return a block of audio data resampled to 16000hz, blocking if necessary."""
        return self.resample(data=self.buffer_queue.get(),
                             input_rate=self.input_rate)

    def read(self):
        """Return a block of audio data, blocking if necessary."""
        return self.buffer_queue.get()

    def destroy(self):
        self.stream.stop_stream()
        self.stream.close()
        self.pa.terminate()

    frame_duration_ms = property(lambda self: 1000 * self.block_size // self.sample_rate)

    def write_wav(self, filename, data):
        logging.warning("write wav %s", filename)
        with wave.open(filename, 'wb') as wf:
            wf.setnchannels(self.CHANNELS)
            # wf.setsampwidth(self.pa.get_sample_size(FORMAT))
            assert self.FORMAT == pyaudio.paInt16
            wf.setsampwidth(2)
            wf.setframerate(self.sample_rate)
            wf.writeframes(data)

    def frame_generator(self):
        """Generator that yields all audio frames from microphone."""
        if self.input_rate == self.RATE_PROCESS:
            while True:
                yield self.read()
        else:
            while True:
                yield self.read_resampled()

def max_decode(output, blank=39):
    # find the argmax of each label at each timestep. the label dimension is reduced.
    pred = np.argmax(output, 1)
    # initialize the sequence as an empty list of the first prediction is blank index
    prev = pred[0]
    seq = [prev] if prev != blank else []
    # iterate through the predictions and condense repeated predictions
    for p in pred[1:]:
        if p != blank and p != prev:
            seq.append(p)
        prev = p

    return seq

if __name__ == '__main__':
    BEAM_WIDTH = 500
    DEFAULT_SAMPLE_RATE = 16000

    import argparse
    parser = argparse.ArgumentParser(description="Stream from microphone to DeepSpeech using VAD")
    parser.add_argument(
        '-w', '--savewav', help="Save .wav files of utterences to given directory"
    )
    parser.add_argument(
        '-f', '--file', help="Read from .wav file instead of microphone"
    )
    parser.add_argument(
        '-md', '--model-dir', help="Path to model directory that contains model, preproc, and config file."
    )
    parser.add_argument(
        '-t', '--tag', type=str, default='', choices=['best', ''], help="tag if 'best' model is desired"
    )
    parser.add_argument(
        '-mn', '--model-name', type=str, default='', help="name of model to override default in get_names method"
    )
    parser.add_argument(
        '-d', '--device', type=int, default=None,
        help="Device input index (Int) as listed by pyaudio.PyAudio.get_device_info_by_index(). If not provided, falls back to PyAudio.get_default_device()."
    )
    parser.add_argument(
        '-r', '--rate', type=int, default=DEFAULT_SAMPLE_RATE,
        help=f"Input device sample rate. Default: {DEFAULT_SAMPLE_RATE}. Your device may require 44100."
    )
    # ctc decoder not currenlty used
    parser.add_argument('-bw', '--beam_width', type=int, default=BEAM_WIDTH,
                        help=f"Beam width used in the CTC decoder when building candidate transcriptions. Default: {BEAM_WIDTH}")

    ARGS = parser.parse_args()
    if ARGS.savewav:
        os.makedirs(ARGS.savewav, exist_ok=True)
    main(ARGS)

ext: py

max_stars_repo_name: dzubke/speech-lite

max_issues_repo_name: dzubke/speech-lite

max_forks_repo_name: dzubke/speech-lite

reason: import numpy,from scipy


text:
import os
import torch
import torchvision
import numpy as np
from pytorch_lightning.callbacks.base import Callback
from import log_histogram_3d

class Comet3DHistogramSampler(Callback):
    def __init__(self, _C,logger):
        super().__init__()
        self.logger = logger

    def on_train_epoch_end(self, trainer, pl_module):
    # def on_after_backward(self, trainer, pl_module):

        model = pl_module.trainer.model
        epoch = model.current_epoch

        if trainer.global_step % 5 == 0:  # don't make the file huge
            params = model.state_dict()
            for k, v in params.items():
                grads = v
                name = f"epoch_{epoch}_params/{k}"
                self.logger.experiment.add_histogram(tag=name, values=grads,
                                                    global_step=self.trainer.global_step)

                name = f"epoch_{epoch}_params/{k}"
                self.logger.experiment.add_histogram(tag=name, values=grads.grad,
                                                    global_step=self.trainer.global_step)

        # Add this to the training loop in the train_one_epoch function
        # for name, param in model.named_parameters():
        #         log_histogram_3d(
        #             values=param,
        #             name=f"epoch_{epoch}_params/{name}",
        #             step=epoch
        #         )

        #         log_histogram_3d(
        #             values=param.grad,
        #             name=f"epoch_{epoch}_grads/{name}",
        #             step=epoch
        #         )

        # For run parameters -
        # Add this to end of epoch loop
    # def on_train_epoch_end(self, trainer, pl_module):
    #     for layer, (name, param) in enumerate(model.named_parameters()):
    #             writer.add_histogram(
    #                 tag=f"run_params/{name}", values=param, global_step=epoch
    #             )

max_stars_repo_name: eungbean/DCGAN-pytorch-lightning-comet

max_issues_repo_name: eungbean/DCGAN-pytorch-lightning-comet

max_forks_repo_name: eungbean/DCGAN-pytorch-lightning-comet

reason: import numpy


text:
# -*- coding: utf-8 -*-
"""
Sample code for generating an airfoil surface, extracting the normal vectors
and visualising it using PyVista.

Based on the following issue: https://github.com/pyvista/pyvista-support/issues/506

Author:
    Fabio Casagrande Hirono
    Sep 2021
"""

import numpy as np
import pyvista as pv

import pywopwop as PWW

# %% *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-
# build a generic airfoil using 'airfoils' package (available on PyPI)
# ---> https://airfoils.readthedocs.io/en/latest/index.html

# from airfoils import Airfoil
# foil = Airfoil.NACA4('4812')
# # foil = Airfoil.NACA4('0012')

# Nx = 102
# x_foil = np.linspace(0, 1, Nx//2)
# y_foil_up = foil.y_upper(x_foil)
# y_foil_lo = foil.y_lower(x_foil)

# # concatenate top and bottom
# x = np.concatenate((x_foil, x_foil[::-1]))
# y = np.concatenate((y_foil_up, y_foil_lo[::-1]))

# *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-
# Alternatively, define data points directly:

x = np.array([ 0.  , 0.02, 0.04, 0.06, 0.08, 0.1 , 0.12, 0.14, 0.16, 0.18, 0.2 ,
                0.22, 0.24, 0.26, 0.28, 0.3 , 0.32, 0.34, 0.36, 0.38, 0.4 , 0.42,
                0.44, 0.46, 0.48, 0.5 , 0.52, 0.54, 0.56, 0.58, 0.6 , 0.62, 0.64,
                0.66, 0.68, 0.7 , 0.72, 0.74, 0.76, 0.78, 0.8 , 0.82, 0.84, 0.86,
                0.88, 0.9 , 0.92, 0.94, 0.96, 0.98, 1.  , 1.  , 0.98, 0.96, 0.94,
                0.92, 0.9 , 0.88, 0.86, 0.84, 0.82, 0.8 , 0.78, 0.76, 0.74, 0.72,
                0.7 , 0.68, 0.66, 0.64, 0.62, 0.6 , 0.58, 0.56, 0.54, 0.52, 0.5 ,
                0.48, 0.46, 0.44, 0.42, 0.4 , 0.38, 0.36, 0.34, 0.32, 0.3 , 0.28,
                0.26, 0.24, 0.22, 0.2 , 0.18, 0.16, 0.14, 0.12, 0.1 , 0.08, 0.06,
                0.04, 0.02, 0.  ])

y = np.array([  0.        ,  0.03805971,  0.05349321,  0.06597081,  0.07664938,
                0.08600053,  0.09427433,  0.10162394,  0.10815206,  0.11393193,
                0.11901819,  0.12345299,  0.1272697 ,  0.13049523,  0.13315164,
                0.13525721,  0.13682722,  0.13787454,  0.13841001,  0.13844282,
                0.13797963,  0.13710874,  0.13595412,  0.13452314,  0.13282156,
                0.13085448,  0.12862644,  0.12614146,  0.12340309,  0.12041439,
                0.11717802,  0.11369624,  0.10997091,  0.10600351,  0.10179516,
                0.09734664,  0.09265833,  0.0877303 ,  0.08256226,  0.07715356,
                0.07150319,  0.06560979,  0.05947162,  0.05308659,  0.04645218,
                0.03956552,  0.03242328,  0.02502174,  0.0173567 ,  0.00942353,
                0.00121706, -0.0012992 ,  0.00117944,  0.00354851,  0.00580537,
                0.00794761,  0.00997303,  0.01187967,  0.01366579,  0.01532991,
                0.01687078,  0.01828745,  0.01957923,  0.0207457 ,  0.02178678,
                0.02270265,  0.02349386,  0.02416127,  0.02470611,  0.02512998,
                0.02543486,  0.02562315,  0.02569767,  0.02566172,  0.02551906,
                0.02527399,  0.02493133,  0.02449652,  0.02397561,  0.02337537,
                0.0227033 ,  0.0219667 ,  0.02103652,  0.01982814,  0.0183556 ,
                0.01663352,  0.01467847,  0.01250933,  0.01014766,  0.00761818,
                0.00494959,  0.00217557, -0.00066369, -0.00351907, -0.00632877,
                -0.00901202, -0.01145764, -0.01350184, -0.01487876, -0.01508707,
                -0.01287165,  0.        ])

Nx = x.shape[0]

# *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-

# z-direction coordinates (span)
Nz = 51
z_max = 5
z = np.linspace(0, z_max, Nz)

# build 3D coordinate array
XYZ = np.zeros((Nx, Nz, 3))
XYZ[:, :, 0] = x[:, np.newaxis]
XYZ[:, :, 1] = y[:, np.newaxis]
XYZ[:, :, 2] = z[np.newaxis, :]

# apply spanwise-varying twist - max 45deg at tip
theta_twist = -np.pi/4

XYZ[:, :, 0] = (np.cos(theta_twist*XYZ[:, :, 2]/z_max)*XYZ[:, :, 0]
                - np.sin(theta_twist*XYZ[:, :, 2]/z_max)*XYZ[:, :, 1])
XYZ[:, :, 1] = (np.sin(theta_twist*XYZ[:, :, 2]/z_max)*XYZ[:, :, 0]
                + np.cos(theta_twist*XYZ[:, :, 2]/z_max)*XYZ[:, :, 1])

# %% *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-
# pass the mesh points to the StructuredGrid constructor, plot mesh points for
# visual inspection

airfoil = pv.StructuredGrid(XYZ[:, :, 0], XYZ[:, :, 1], XYZ[:, :, 2])

airfoil.plot(background='w', show_edges=True, pbr=True, metallic=1.0, roughness=0.4)

# %% *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-
# Use pyvista to obtain surface normal vectors, plot normal vectors for
# inspection

# extract PolyData representation of mesh
airfoil2 = airfoil.extract_surface()

# extract normal vectors as numpy array
airfoil2 = airfoil2.compute_normals(flip_normals=True)
normals = (airfoil2.point_normals).reshape(XYZ.shape)

# plot mesh with normal vectors and save a screenshot
plotter = pv.Plotter()
plotter.add_mesh(airfoil2, show_edges=True)
plotter.add_arrows(airfoil2.points, airfoil2.point_normals, mag=0.02)

plotter.show_bounds(bounds=[0., 1, 0., 1, 0, 5], grid='front', location='outer',
                    all_edges=True)
plotter.show(auto_close=False)
# plotter.show(screenshot='my_normal_vectors.png')

# %%
# Build pywopwop patch object with structured grid geometry

# create empty instance of PWW Patch File
rotorblade_patch = PWW.PWWPatch()

# add header info
rotorblade_patch.geom_type = 'geometry'
rotorblade_patch.is_structured = True
rotorblade_patch.geometry_time_type = 'constant'
rotorblade_patch.centered_type = 'node'
rotorblade_patch.float_type = 'single'
rotorblade_patch.iblank_included = 'n'
rotorblade_patch.units_string = 'Pa'

rotorblade_patch.geometry_comment = 'Test geometry file for single blade'

# add rotor blade geometry and normals to patch obj
rotorblade_patch.add_StructuredZone('Rotor blade', XYZ.T, normals.T,
                                    calc_thickness_noise=True)

# print info
rotorblade_patch.print_info()

# create PSU-WOPWOP patch file
rotorblade_geom_filename = 'NACA4812_patch.dat'
rotorblade_patch.write_geometry_file(rotorblade_geom_filename)

max_stars_repo_name: fchirono/pywopwop

max_issues_repo_name: fchirono/pywopwop

max_forks_repo_name: fchirono/pywopwop