text:
import os
import time

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import AutoConfig, AutoModel

from . import char_lstm
from . import decode_chart
from . import nkutil
from .partitioned_transformer import (
    ConcatPositionalEncoding,
    FeatureDropout,
    PartitionedTransformerEncoder,
    PartitionedTransformerEncoderLayer,
)
from . import parse_base
from . import retokenization
from . import subbatching
from . import tetra_tag
from . import vector_quantize

class ChartParser(nn.Module, parse_base.BaseParser):
    def __init__(
        self,
        tag_vocab,
        label_vocab,
        char_vocab,
        hparams,
        pretrained_model_path=None,
    ):
        super().__init__()
        self.config = locals()
        self.config.pop("self")
        self.config.pop("__class__")
        self.config.pop("pretrained_model_path")
        self.config["hparams"] = hparams.to_dict()

        self.tag_vocab = tag_vocab
        self.label_vocab = label_vocab
        self.char_vocab = char_vocab

        self.d_model = hparams.d_model
        self.pretrained_divide = hparams.pretrained_divide

        self.two_label_subspan = hparams.two_label_subspan
        try:
            self.uni = hparams.uni
        except:
            self.uni = False
        try:
            self.all_layers_uni = hparams.all_layers_uni
        except:
            self.all_layers_uni = False

        self.d_cats = hparams.discrete_cats
        try:
            self.use_vq = hparams.use_vq
        except:
            self.use_vq = False

        try:
            self.tags_per_word = hparams.tags_per_word
            assert self.d_cats % self.tags_per_word == 0, "tags-per-word must divide d_cats"
        except:
            self.tags_per_word = 1

        try:
            self.mask = hparams.mask
        except:
            self.mask = (False for _ in range(self.d_cats))

        self.back_cycle = None

        self.char_encoder = None
        self.pretrained_model = None

        # try:
        if hasattr(hparams, 'use_w2v_clustered_lexicon') and hparams.use_w2v_clustered_lexicon:
            from gensim.models import KeyedVectors
            print('Using Word2Vec')
            clustered_lexicon = KeyedVectors.load_word2vec_format(hparams.use_w2v_clustered_lexicon)

            def get_vec(x):
                if x in clustered_lexicon:
                    return clustered_lexicon[x]
                return clustered_lexicon['UNK']
            self.clustered_lexicon = get_vec
            print('Loaded Lexicon')

        elif hasattr(hparams, 'use_ft_clustered_lexicon') and hparams.use_ft_clustered_lexicon:
            print('Using FastText')
            from gensim.models.fasttext import load_facebook_model
            clustered_lexicon = load_facebook_model(hparams.use_ft_clustered_lexicon)
            self.clustered_lexicon = lambda x: clustered_lexicon.wv[x]
            print('Loaded Lexicon')
        else:
            self.clustered_lexicon = None

        # except:
        #     print('Error loading Lexicon')
        #     self.clustered_lexicon = None

        if hparams.use_chars_lstm:
            assert (
                not hparams.use_pretrained
            ), "use_chars_lstm and use_pretrained are mutually exclusive"
            self.retokenizer = char_lstm.RetokenizerForCharLSTM(
                self.char_vocab)
            self.char_encoder = char_lstm.CharacterLSTM(
                max(self.char_vocab.values()) + 1,
                hparams.d_char_emb,
                hparams.d_model // 2,  # Half-size to leave room for
                # partitioned positional encoding
                char_dropout=hparams.char_lstm_input_dropout,
            )
        elif hparams.use_pretrained or self.clustered_lexicon:
            if hparams.use_pretrained:
                if pretrained_model_path is None:
                    self.retokenizer = retokenization.Retokenizer(
                        hparams.pretrained_model, retain_start_stop=True
                    )
                    self.pretrained_model = AutoModel.from_pretrained(
                        hparams.pretrained_model
                    )
                else:
                    self.retokenizer = retokenization.Retokenizer(
                        pretrained_model_path, retain_start_stop=True
                    )
                    self.pretrained_model = AutoModel.from_config(
                        AutoConfig.from_pretrained(pretrained_model_path)
                    )
                d_pretrained = self.pretrained_model.config.hidden_size
            self.use_forced_lm = hparams.use_forced_lm
            self.bpe_dropout = hparams.bpe_dropout

            if self.clustered_lexicon:
                d_pretrained = 300
            if hparams.use_encoder:
                if self.d_cats > 0 and self.use_vq:
                    self.project_pretrained = nn.Linear(
                        d_pretrained, hparams.d_model // 2, bias=False
                    )
                    self.vq = vector_quantize.VectorQuantize(
                        dim=hparams.d_model // 2 if not self.clustered_lexicon else 300,
                        n_embed=self.d_cats,
                        decay=hparams.vq_decay,
                        commitment=hparams.vq_commitment,
                        wait_steps=hparams.vq_wait_steps,
                        observe_steps=hparams.vq_observe_steps,
                        coreset_size_multiplier=hparams.vq_coreset_size_multiplier,
                    )
                    self.commit_loss_accum = 0.0
                elif self.d_cats > 0:
                    self.project_in = nn.Linear(
                        d_pretrained, self.d_cats, bias=False)

                    # scaling to adjust output of gpt2
                    self.project_in.weight.data *= 1e-3

                    self.project_out = nn.Linear(
                        self.d_cats, hparams.d_model // 2, bias=False)
                else:
                    self.project_pretrained = nn.Linear(
                        d_pretrained, hparams.d_model // 2, bias=False
                    )
            else:
                self.project_pretrained = nn.Linear(
                    d_pretrained, hparams.d_model, bias=False
                )

        if hparams.use_encoder:
            self.morpho_emb_dropout = FeatureDropout(
                hparams.morpho_emb_dropout)
            self.add_timing = ConcatPositionalEncoding(
                d_model=hparams.d_model,
                max_len=hparams.encoder_max_len,
            )
            encoder_layer = PartitionedTransformerEncoderLayer(
                hparams.d_model,
                n_head=hparams.num_heads,
                d_qkv=hparams.d_kv,
                d_ff=hparams.d_ff,
                ff_dropout=hparams.relu_dropout,
                residual_dropout=hparams.residual_dropout,
                attention_dropout=hparams.attention_dropout,
            )
            first_layer = None
            try:
                if hparams.first_heads > 0:
                    first_layer = PartitionedTransformerEncoderLayer(
                        hparams.d_model,
                        n_head=hparams.first_heads,
                        d_qkv=hparams.d_kv,
                        d_ff=hparams.d_ff,
                        ff_dropout=hparams.relu_dropout,
                        residual_dropout=hparams.residual_dropout,
                        attention_dropout=hparams.attention_dropout,
                    )
            except:
                hparams.first_heads = -1
            self.encoder_gum = hparams.encoder_gum
            if self.encoder_gum:
                self.project_pretrained = nn.Linear(
                    d_pretrained, hparams.d_model // 2, bias=False
                )
                num_heads = hparams.first_heads if hparams.first_heads > 0 else hparams.num_heads
                self.w_qkv_c = nn.Parameter(torch.Tensor(
                    num_heads, self.d_model // 2, hparams.d_kv // 2))
                self.w_qkv_p = nn.Parameter(torch.Tensor(
                    num_heads, self.d_model // 2, hparams.d_kv // 2))
            self.encoder = PartitionedTransformerEncoder(
                encoder_layer, hparams.num_layers, first_layer
            )
            try:
                if hparams.back_cycle:
                    self.back_add_timing = ConcatPositionalEncoding(
                        d_model=hparams.d_model,
                        max_len=hparams.encoder_max_len,
                    )
                    self.back_use_gold_trees = hparams.back_use_gold_trees
                    self.back_cycle = PartitionedTransformerEncoder(
                        encoder_layer, hparams.back_layers
                    )
                    if hparams.two_label or hparams.two_label_subspan:
                        self.tetra_tag_system = tetra_tag.TetraTagSystem(
                            tag_vocab=['L/S', 'R/S', 'l', 'r'])
                        self.tetra_leaves = [2, 3]
                    else:
                        config = AutoConfig.from_pretrained(
                            'kitaev/tetra-tag-en')
                        tag_vocab = [config.id2label[i]
                                     for i in sorted(config.id2label.keys())]
                        self.tetra_tag_system = tetra_tag.TetraTagSystem(
                            tag_vocab=tag_vocab)
                        self.tetra_leaves = [i for i in range(
                            len(tag_vocab)) if tag_vocab[i][0] in 'lr']

                    self.back_project = nn.Linear(
                        len(self.tetra_tag_system.tag_vocab) + 1, hparams.d_model // 2, bias=False
                    )
                    self.back_loss_constant = hparams.back_loss_constant
                    if self.d_cats > 0:
                        seq_dim = self.d_cats
                    else:
                        seq_dim = hparams.d_model // 2
                    self.f_back = nn.Sequential(
                        nn.Linear(hparams.d_model, hparams.d_label_hidden),
                        nn.LayerNorm(hparams.d_label_hidden),
                        nn.ReLU(),
                        nn.Linear(hparams.d_label_hidden, seq_dim),
                    )
                    if hparams.back_loss_type == 'ce':
                        from .loss_functions import cel
                        self.back_criterion = cel
                    elif hparams.back_loss_type == 'tvd':
                        from .loss_functions import tvd
                        self.back_criterion = tvd
                    elif hparams.back_loss_type == 'kl':
                        from .loss_functions import kl
                        self.back_criterion = kl(
                            reduction='sum')
                    elif hparams.back_loss_type == 'js':
                        from .loss_functions import js_gen
                        self.back_criterion = js_gen(
                            reduction='sum')
                    elif hparams.back_loss_type == 'emd':
                        from .loss_functions import emd
                        self.back_criterion = emd
                    elif hparams.back_loss_type == 'mse':
                        self.back_criterion = nn.MSELoss(reduction='sum')
                    else:
                        raise ValueError(
                            "Invalid type of loss for backwards cycle")
                    #  measures of simularities between probability distributions (wasserstein, earth movers, section 2 of wasserstein gan paper)
                    # loss from distribution difference (softmax can end up w inf)

            except:
                self.back_cycle = None
        else:
            self.morpho_emb_dropout = None
            self.add_timing = None
            self.encoder = None

        self.f_label = nn.Sequential(
            nn.Linear(hparams.d_model, hparams.d_label_hidden),
            nn.LayerNorm(hparams.d_label_hidden),
            nn.ReLU(),
            nn.Linear(hparams.d_label_hidden, max(label_vocab.values())),
        )

        if hparams.predict_tags:
            self.f_tag = nn.Sequential(
                nn.Linear(hparams.d_model, hparams.d_tag_hidden),
                nn.LayerNorm(hparams.d_tag_hidden),
                nn.ReLU(),
                nn.Linear(hparams.d_tag_hidden, max(tag_vocab.values()) + 1),
            )
            self.tag_loss_scale = hparams.tag_loss_scale
            self.tag_from_index = {i: label for label, i in tag_vocab.items()}
        else:
            self.f_tag = None
            self.tag_from_index = None

        self.decoder = decode_chart.ChartDecoder(
            label_vocab=self.label_vocab,
            force_root_constituent=hparams.force_root_constituent,
        )
        self.criterion = decode_chart.SpanClassificationMarginLoss(
            reduction="sum", force_root_constituent=hparams.force_root_constituent
        )

        self.parallelized_devices = None

    def set_mask(self, m):
        assert len(
            m) == self.d_cats, "Length of given mask does not match number of dimensions"
        self.mask = tuple(b for b in m)

        self.config["hparams"]['mask'] = self.mask

    @property
    def device(self):
        if self.parallelized_devices is not None:
            return self.parallelized_devices[0]
        else:
            return next(self.f_label.parameters()).device

    @property
    def output_device(self):
        if self.parallelized_devices is not None:
            return self.parallelized_devices[1]
        else:
            return next(self.f_label.parameters()).device

    def parallelize(self, *args, **kwargs):
        self.parallelized_devices = (torch.device(
            "cuda", 0), torch.device("cuda", 1))
        for child in self.children():
            if child != self.pretrained_model:
                child.to(self.output_device)

        self.pretrained_model.parallelize(*args, **kwargs)

    @classmethod
    def from_trained(cls, model_path):
        if os.path.isdir(model_path):
            # Multi-file format used when exporting models for release.
            # Unlike the checkpoints saved during training, these files include
            # all tokenizer parameters and a copy of the pre-trained model
            # config (rather than downloading these on-demand).
            config = AutoConfig.from_pretrained(model_path).benepar
            state_dict = torch.load(
                os.path.join(model_path, "benepar_model.bin"), map_location="cpu"
            )
            config["pretrained_model_path"] = model_path
        else:
            # Single-file format used for saving checkpoints during training.
            data = torch.load(model_path, map_location="cpu")
            config = data["config"]
            state_dict = data["state_dict"]

        hparams = config["hparams"]

        if "force_root_constituent" not in hparams:
            hparams["force_root_constituent"] = True
        if "bpe_dropout" not in hparams:
            hparams["bpe_dropout"] = 0.0

        config["hparams"] = nkutil.HParams(**hparams)
        parser = cls(**config)
        parser.load_state_dict(state_dict, strict=False)

        if hasattr(parser, 'vq'):
            # HACK: Skip vq initialization when restoring from checkpoint.
            parser.vq.wait_steps_remaining = 0
            parser.vq.observe_steps_remaining = 0

        return parser

    def encode(self, example, use_bpe_dropout=False):

        if self.two_label_subspan is not None and self.two_label_subspan is not False:
            example = self.two_label_subspan(example, self.tree_transform)

        if self.char_encoder is not None:
            encoded = self.retokenizer(example.words, return_tensors="np")
        else:
            encoded = self.retokenizer(example.words, example.space_after,
                                       dropout=(self.bpe_dropout if use_bpe_dropout else None))

        if self.clustered_lexicon:
            wvs = [self.clustered_lexicon('X').copy()] + \
                [self.clustered_lexicon(
                    x).copy() for x in example.words] + [self.clustered_lexicon('X').copy()]
            encoded['wv'] = torch.FloatTensor(wvs)
            encoded['wv_mask'] = torch.ones(
                len(wvs), dtype=torch.bool)
            encoded['wv_mask'][[0, -1]] = False

        if example.tree is not None:
            encoded["span_labels"] = torch.tensor(
                self.decoder.chart_from_tree(example.tree)
            )
            if self.f_tag is not None:
                encoded["tag_labels"] = torch.tensor(
                    [-100] + [self.tag_vocab[tag]
                              for _, tag in example.pos()] + [-100]
                )
            if self.back_cycle is not None and self.back_use_gold_trees:
                tags = [len(self.tetra_tag_system.tag_vocab)] + self.tetra_tag_system.ids_from_tree(
                    example.tree) + [len(self.tetra_tag_system.tag_vocab)]
                encoded['tetra_tags'] = torch.tensor(tags)
                mask = [i in self.tetra_leaves for i in tags]
                mask[0] = mask[-1] = True
                encoded['tetra_tags_mask'] = torch.tensor(mask)
                padding_mask = [True for i in tags]
                encoded['tetra_padding_mask'] = torch.tensor(padding_mask)
        return encoded

    def pad_encoded(self, encoded_batch):
        batch = self.retokenizer.pad(
            [
                {
                    k: v
                    for k, v in example.items()
                    if (k not in ["span_labels", "tag_labels", "tetra_tags",
                                  "tetra_tags_mask", "tetra_padding_mask", 'wv', 'wv_mask'])
                }
                for example in encoded_batch
            ],
            return_tensors="pt",
        )
        if encoded_batch and "span_labels" in encoded_batch[0]:
            batch["span_labels"] = decode_chart.pad_charts(
                [example["span_labels"] for example in encoded_batch]
            )
        if encoded_batch and "tag_labels" in encoded_batch[0]:
            batch["tag_labels"] = nn.utils.rnn.pad_sequence(
                [example["tag_labels"] for example in encoded_batch],
                batch_first=True,
                padding_value=-100,
            )
        if encoded_batch and "wv" in encoded_batch[0]:
            batch["wv"] = nn.utils.rnn.pad_sequence(
                [example["wv"] for example in encoded_batch],
                batch_first=True
            )
        if encoded_batch and "wv_mask" in encoded_batch[0]:
            batch["wv_mask"] = nn.utils.rnn.pad_sequence(
                [example["wv_mask"] for example in encoded_batch],
                batch_first=True
            )
        if encoded_batch and "tetra_tags" in encoded_batch[0]:
            batch["tetra_tags"] = nn.utils.rnn.pad_sequence(
                [example["tetra_tags"] for example in encoded_batch],
                batch_first=True,
                padding_value=len(self.tetra_tag_system.tag_vocab)
            )
        if encoded_batch and "tetra_padding_mask" in encoded_batch[0]:
            batch["tetra_padding_mask"] = nn.utils.rnn.pad_sequence(
                [example["tetra_padding_mask"] for example in encoded_batch],
                batch_first=True
            )
        if encoded_batch and "tetra_tags_mask" in encoded_batch[0]:
            batch["tetra_tags_mask"] = nn.utils.rnn.pad_sequence(
                [example["tetra_tags_mask"] for example in encoded_batch],
                batch_first=True
            )
        return batch

    def _get_lens(self, encoded_batch):
        if self.pretrained_model is not None:
            return [len(encoded["input_ids"]) for encoded in encoded_batch]
        return [len(encoded["valid_token_mask"]) for encoded in encoded_batch]

    def encode_and_collate_subbatches(self, examples, subbatch_max_tokens):
        batch_size = len(examples)
        batch_num_tokens = sum(len(x.words) for x in examples)
        encoded = [self.encode(example, use_bpe_dropout=True)
                   for example in examples]

        res = []
        for ids, subbatch_encoded in subbatching.split(
            encoded, costs=self._get_lens(encoded), max_cost=subbatch_max_tokens
        ):
            subbatch = self.pad_encoded(subbatch_encoded)
            subbatch["batch_size"] = batch_size
            subbatch["batch_num_tokens"] = batch_num_tokens
            res.append((len(ids), subbatch))
        return res

    def forward(self, batch, tau=1.0, en_tau=0.0, return_cat_logits=False, force_cats=None):
        category_ret = None
        commit_loss = None
        if force_cats is None:
            # Begin calculating from batch
            valid_token_mask = batch["valid_token_mask"].to(self.output_device)

            if (
                self.encoder is not None
                and valid_token_mask.shape[1] > self.add_timing.timing_table.shape[0]
            ):
                raise ValueError(
                    "Sentence of length {} exceeds the maximum supported length of "
                    "{}".format(
                        valid_token_mask.shape[1] - 2,
                        self.add_timing.timing_table.shape[0] - 2,
                    )
                )
            if self.char_encoder is not None:
                assert isinstance(self.char_encoder, char_lstm.CharacterLSTM)
                char_ids = batch["char_ids"].to(self.device)
                extra_content_annotations = self.char_encoder(
                    char_ids, valid_token_mask)
            elif self.pretrained_model is not None or 'wv' in batch:
                if 'wv' in batch:
                    # using clustered lexicon
                    features = batch['wv'].to(
                        self.output_device)
                    valid_token_mask = batch['wv_mask'].to(
                        self.output_device)
                else:
                    assert not self.clustered_lexicon
                    input_ids = batch["input_ids"].to(self.device)
                    words_from_tokens = batch["words_from_tokens"].to(
                        self.output_device)
                    pretrained_attention_mask = batch["attention_mask"].to(
                        self.device)

                    extra_kwargs = {}
                    if "token_type_ids" in batch:
                        extra_kwargs["token_type_ids"] = batch["token_type_ids"].to(
                            self.device)
                    if "decoder_input_ids" in batch:
                        extra_kwargs["decoder_input_ids"] = batch["decoder_input_ids"].to(
                            self.device
                        )
                        extra_kwargs["decoder_attention_mask"] = batch[
                            "decoder_attention_mask"
                        ].to(self.device)
                    if self.retokenizer.is_t5 and self.use_forced_lm:
                        pretrained_out = self.pretrained_model(
                            input_ids=input_ids[:, :1],
                            attention_mask=pretrained_attention_mask[:, :1],
                            return_dict=True,
                            **extra_kwargs
                        )
                    else:
                        pretrained_out = self.pretrained_model(
                            input_ids, attention_mask=pretrained_attention_mask, return_dict=True, **extra_kwargs
                        )
                    features = pretrained_out.last_hidden_state.to(
                        self.output_device) / self.pretrained_divide
                    features = features[
                        torch.arange(features.shape[0])[:, None],
                        # Note that words_from_tokens uses index -100 for invalid positions
                        F.relu(words_from_tokens),
                    ]
                    # print('Using Features', features)
                    assert not self.clustered_lexicon

                features.masked_fill_(~valid_token_mask[:, :, None], 0)

                if self.encoder is not None:
                    if self.d_cats > 0 and self.use_vq:
                        # Create a mask that excludes start and stop tokens.
                        # This means that discrete category decisions are made
                        # for each word in the sentence, but not for any special
                        # token. The encoder transformer layers have start/stop
                        # tokens, but those will receive an all-zero content
                        # embedding. Note that GPT-2 does not use a start token,
                        # so the default setting in retokenization.py is to use
                        # the GPT-2 *stop token* embedding at that position.
                        # That would violate incrementality, so we mask it out.
                        quantization_mask = valid_token_mask.clone()
                        quantization_mask[:, 0] = False
                        quantization_mask[
                            torch.arange(features.shape[0]),
                            valid_token_mask.sum(1) - 1,
                        ] = False

                        projected_features = self.project_pretrained(features)
                        unquantized_features = projected_features[quantization_mask]

                        (quantized_features, categories, commit_loss, dist) = self.vq(
                            unquantized_features,
                            batch["batch_num_tokens"] if "batch_num_tokens" in batch else None
                        )

                        if tau > 0.0:
                            assert self.training, "expected no annealing during eval"
                            quantized_features = (
                                (1.0 - tau) * quantized_features
                                + tau * unquantized_features)

                        extra_content_annotations = torch.zeros_like(
                            projected_features)
                        extra_content_annotations[quantization_mask] = quantized_features

                        if return_cat_logits:
                            if dist is None:
                                category_ret = None
                            else:
                                category_ret = torch.zeros(
                                    (*
                                     projected_features.shape[:-1], self.d_cats),
                                    dtype=projected_features.dtype,
                                    device=categories.device,
                                )
                                category_ret[quantization_mask] = dist
                        else:
                            category_ret = torch.zeros(
                                projected_features.shape[:-1],
                                dtype=categories.dtype,
                                device=categories.device,
                            )
                            category_ret[quantization_mask] = categories

                        if tau > 0.0:
                            assert self.training, "expected no annealing during eval"
                            extra_content_annotations[:,
                                                      0] = tau * projected_features[:, 0]
                            extra_content_annotations[
                                torch.arange(features.shape[0]),
                                valid_token_mask.sum(1) - 1,
                            ] = tau * projected_features[
                                torch.arange(features.shape[0]),
                                valid_token_mask.sum(1) - 1,
                            ]

                            commit_loss = commit_loss * (1.0 - tau)

                        commit_loss = commit_loss * \
                            unquantized_features.shape[0] * \
                            unquantized_features.shape[1]
                    elif self.d_cats > 0:
                        assert self.tags_per_word == 1, 'Disabling tags per word'

                        category_logits = self.project_in(features)

                        mask = torch.BoolTensor(
                            np.tile(self.mask, (category_logits.shape[0], category_logits.shape[1], 1))).to(self.device)
                        category_logits = category_logits.masked_fill(
                            mask, -1e9)
                        b, w, d = category_logits.shape
                        if tau > 0:
                            cats = F.gumbel_softmax(
                                category_logits, dim=-1, tau=tau, hard=True)
                        else:
                            # when tau == 0, take the argmax (for deterministic testing)
                            max_idx = torch.argmax(category_logits, dim=-1)
                            cats = F.one_hot(max_idx, num_classes=self.d_cats).type(
                                torch.FloatTensor).to(self.output_device)

                        if return_cat_logits:
                            category_ret = category_logits  # .reshape(b, w, d)
                        else:
                            category_ret = cats  # .reshape(b, w, d)

                        # cats = cats.reshape(b, w, d)

                        extra_content_annotations = self.project_out(cats)
                    else:
                        extra_content_annotations = self.project_pretrained(
                            features)
                        category_ret = extra_content_annotations
            # end calculating from batch
        elif self.use_vq:
            # Begin forcing vq categories
            assert self.encoder is not None and self.d_cats > 0, "Forcing categories only supported with discretization"
            batch_size = len(force_cats)
            max_len = max([len(x) for x in force_cats]) + \
                2  # +2 for start/stop

            valid_token_mask = torch.zeros(
                (batch_size, max_len), dtype=torch.bool)
            quantization_mask = valid_token_mask.clone()
            categories = []

            for i, sent_categories in enumerate(force_cats):
                categories.extend(sent_categories)
                valid_token_mask[i, :len(sent_categories) + 2] = True
                quantization_mask[i, 1:len(sent_categories) + 1] = True

            categories = torch.tensor(
                categories, dtype=torch.long, device=self.output_device)
            valid_token_mask = valid_token_mask.to(self.output_device)
            quantization_mask = quantization_mask.to(self.output_device)

            quantized_features = F.embedding(
                categories, self.vq.embed.transpose(0, 1))

            extra_content_annotations = torch.zeros(
                (batch_size, max_len, quantized_features.shape[-1]),
                dtype=quantized_features.dtype, device=quantized_features.device
            )
            extra_content_annotations[quantization_mask] = quantized_features
        else:
            # Begin forcing gumbel categories
            assert self.encoder is not None and self.d_cats > 0, "Forcing categories only supported with discretization"
            assert not self.use_vq, "Not implemented with vector quantization"
            category_ret = force_cats
            force_cats = torch.LongTensor(force_cats)
            force_cats = force_cats.reshape(
                force_cats.shape[0], force_cats.shape[1], self.tags_per_word)
            cats = F.one_hot(force_cats,
                             num_classes=self.d_cats // self.tags_per_word).to(self.output_device).float()
            b, w, tpw, d = cats.shape
            cats = cats.reshape(b, w, tpw * d)
            valid_token_mask = np.ones((b, w))

            valid_token_mask = torch.BoolTensor(
                valid_token_mask).to(self.output_device)
            extra_content_annotations = self.project_out(cats)

        if self.encoder is not None:
            uni_mask = None
            if self.all_layers_uni or self.uni:
                # converting to backwards masks
                b, t = valid_token_mask.shape
                uni_mask = torch.tril(
                    valid_token_mask.repeat(1, t).reshape(b, t, t))
                if self.all_layers_uni:
                    valid_token_mask, uni_mask = uni_mask, None

                # set mask to b x t x t

            if self.d_cats > 0 and not self.use_vq and tau == 0:
                # disabling dropout to get deterministic results for analysis
                # TODO is this step needed?
                encoder_in = self.add_timing(extra_content_annotations)
            else:
                encoder_in = self.add_timing(
                    self.morpho_emb_dropout(extra_content_annotations)
                )
            if self.encoder_gum:
                assert self.d_cats > 0
                pretrained_encoder_data = self.add_timing(
                    self.morpho_emb_dropout(self.project_pretrained(features)))
                v_c, v_p = torch.chunk(encoder_in, 2, dim=-1)
                v_c = torch.einsum("btf,hfa->bhta", v_c, self.w_qkv_c)
                v_p = torch.einsum("btf,hfa->bhta", v_p, self.w_qkv_p)

                v_cp = torch.cat([v_c, v_p], dim=-1)

                annotations = self.encoder(
                    pretrained_encoder_data, valid_token_mask,
                    att_v=v_cp, in_x=encoder_in, first_layer_mask=uni_mask, tau=en_tau)
                del v_c
                del v_p
                del v_cp
            else:
                annotations = self.encoder(
                    encoder_in, valid_token_mask, first_layer_mask=uni_mask, tau=en_tau)
            # Rearrange the annotations to ensure that the transition to
            # fenceposts captures an even split between position and content.
            # TODO(nikita): try alternatives, such as omitting position entirely
            annotations = torch.cat(
                [
                    annotations[..., 0::2],
                    annotations[..., 1::2],
                ],
                -1,
            )
        else:
            assert self.pretrained_model is not None
            annotations = self.project_pretrained(features)

        if self.f_tag is not None:
            tag_scores = self.f_tag(annotations)
        else:
            tag_scores = None

        fencepost_annotations = torch.cat(
            [
                annotations[:, :-1, : self.d_model // 2],
                annotations[:, 1:, self.d_model // 2:],
            ],
            -1,
        )

        # Note that the bias added to the final layer norm is useless because
        # this subtraction gets rid of it
        span_features = (
            torch.unsqueeze(fencepost_annotations, 1)
            - torch.unsqueeze(fencepost_annotations, 2)
        )[:, :-1, 1:]
        span_scores = self.f_label(span_features)
        span_scores = torch.cat(
            [span_scores.new_zeros(
                span_scores.shape[:-1] + (1,)), span_scores], -1
        )
        return span_scores, tag_scores, category_ret, commit_loss

    def split_tag(self, original_idx, copy_idx):

        # give this tag the same probability of being chosen as original
        self.project_in.weight.data[copy_idx] = self.project_in.weight.data[original_idx]

        # give this tag the same features for the end transformer
        self.project_out.weight.data[:,
                                     copy_idx] = self.project_out.weight.data[:, original_idx]

    def compute_loss(self, batch, tau=1.0, en_tau=1.0):
        span_scores, tag_scores, cats, commit_loss = self.forward(
            batch, tau=tau, en_tau=en_tau, return_cat_logits=True)

        span_labels = batch["span_labels"].to(span_scores.device)
        span_loss = self.criterion(span_scores, span_labels)

        span_loss = span_loss / batch["batch_size"]

        if self.use_vq:
            commit_loss = commit_loss / batch["batch_num_tokens"]
        else:
            commit_loss = 0.0

        if self.back_cycle and cats is not None:
            if self.back_use_gold_trees:
                if self.d_cats > 0:
                    back_loss = self.tetra_to_tags(batch, cats)
                else:
                    back_loss = self.tetra_to_annotations(batch, cats)
            else:
                back_loss = self.predicted_tetra_to_tags(
                    batch, cats, span_scores)

            span_loss += back_loss / batch["batch_num_tokens"]

        if self.use_vq:
            self.commit_loss_accum += float(commit_loss.cpu())

        if tag_scores is None:
            return span_loss + commit_loss
        else:
            tag_labels = batch["tag_labels"].to(tag_scores.device)
            tag_loss = self.tag_loss_scale * F.cross_entropy(
                tag_scores.reshape((-1, tag_scores.shape[-1])),
                tag_labels.reshape((-1,)),
                reduction="sum",
                ignore_index=-100,
            )
            tag_loss = tag_loss / batch["batch_num_tokens"]
            return span_loss + tag_loss + commit_loss

    def tetra_to_tags(self, batch, cats):
        assert self.back_cycle is not None and self.back_use_gold_trees

        # tree to tags
        encoder_in = F.one_hot(batch["tetra_tags"].to(self.device), num_classes=len(
            self.tetra_tag_system.tag_vocab) + 1).float()
        encoder_in = self.back_project(encoder_in)
        encoder_in = self.back_add_timing(self.morpho_emb_dropout(encoder_in))

        annotations = self.back_cycle(
            encoder_in, batch['tetra_padding_mask'].to(self.device))

        logits = self.f_back(annotations)
        loss = self.back_criterion(
            logits[batch["tetra_tags_mask"].to(self.device)], cats[batch["valid_token_mask"].to(self.device)]) / 2

        # back prop experiment
        loss += self.back_criterion(
            cats[batch["valid_token_mask"].to(self.device)], logits[batch["tetra_tags_mask"].to(self.device)]) / 2

        # for later -- use this as a prior for inference
        # if np.random.random() < 0.001:
        #     print('back loss: ', loss.data.cpu().numpy())

        return self.back_loss_constant * loss

    def tetra_to_annotations(self, batch, eca):
        assert self.back_cycle is not None and self.back_use_gold_trees

        # tree to tags
        encoder_in = F.one_hot(batch["tetra_tags"].to(self.device), num_classes=len(
            self.tetra_tag_system.tag_vocab) + 1).float()
        encoder_in = self.back_project(encoder_in)
        encoder_in = self.back_add_timing(self.morpho_emb_dropout(encoder_in))

        annotations = self.back_cycle(
            encoder_in, batch['tetra_padding_mask'].to(self.device))
        logits = self.f_back(annotations)
        loss = self.back_criterion(
            logits[batch["tetra_tags_mask"].to(self.device)], eca[batch["valid_token_mask"].to(self.device)]) / 2

        # back prop experiment
        loss += self.back_criterion(
            eca[batch["valid_token_mask"].to(self.device)], logits[batch["tetra_tags_mask"].to(self.device)]) / 2

        # for later -- use this as a prior for inference
        if np.random.random() < 0.001:
            print('back loss: ', loss.data.cpu().numpy())

        return self.back_loss_constant * loss

    def _parse_encoded(
        self, examples, encoded, return_compressed=False, return_scores=False, return_cats=False, tau=0.0, en_tau=0.0
    ):
        with torch.no_grad():
            if self.check_force_cats(examples):
                span_scores, tag_scores, categories, _ = self.forward(
                    batch=None, force_cats=examples)
                if self.use_vq:
                    lengths = np.array([len(example) for example in examples])
                else:
                    lengths = np.array(
                        [len(example) - 2 for example in examples])
            else:
                batch = self.pad_encoded(encoded)
                span_scores, tag_scores, categories, _ = self.forward(
                    batch, tau)
                lengths = batch["valid_token_mask"].sum(-1) - 2
                lengths = lengths.to(span_scores.device)

            if return_scores:
                span_scores_np = span_scores.cpu().numpy()
            else:
                # Start/stop tokens don't count, so subtract 2
                charts_np = self.decoder.charts_from_pytorch_scores_batched(
                    span_scores, lengths
                )
            if tag_scores is not None:
                tag_ids_np = tag_scores.argmax(-1).cpu().numpy()
            else:
                tag_ids_np = None

        if self.check_force_cats(examples):
            for i in range(len(examples)):
                if self.use_vq:
                    yield self.decoder.tree_from_chart(charts_np[i], leaves=[str(cat) for cat in examples[i]])
                else:
                    yield self.decoder.tree_from_chart(charts_np[i], leaves=[str(cat) for cat in examples[i][1:-1]])
            return

        categories_np = categories.cpu().numpy()

        for i in range(len(encoded)):
            example_len = len(examples[i].words)
            if return_scores:
                yield span_scores_np[i, :example_len, :example_len]
            elif return_compressed:
                output = self.decoder.compressed_output_from_chart(
                    charts_np[i])
                if tag_ids_np is not None:
                    output = output.with_tags(
                        tag_ids_np[i, 1: example_len + 1])
                yield output
            else:
                if tag_scores is None:
                    leaves = examples[i].pos()
                else:
                    predicted_tags = [
                        self.tag_from_index[i]
                        for i in tag_ids_np[i, 1: example_len + 1]
                    ]
                    leaves = [
                        (word, predicted_tag)
                        for predicted_tag, (word, gold_tag) in zip(
                            predicted_tags, examples[i].pos()
                        )
                    ]
                if return_cats and self.use_vq:
                    yield self.decoder.tree_from_chart(charts_np[i], leaves=leaves), categories_np[i, 1:example_len + 1]
                elif return_cats:
                    yield self.decoder.tree_from_chart(charts_np[i], leaves=leaves), categories[i]
                else:
                    yield self.decoder.tree_from_chart(charts_np[i], leaves=leaves)

    def check_force_cats(self, examples):
        return self.d_cats > 0 and (examples is None or isinstance(examples[0], (list, np.ndarray)))

    def parse(
        self,
        examples,
        return_compressed=False,
        return_scores=False,
        subbatch_max_tokens=None,
        return_cats=False,
        tau=0.0,
        en_tau=0.0,
        return_encoded=False
    ):
        training = self.training
        self.eval()
        if self.check_force_cats(examples):
            encoded = None
        else:
            encoded = [self.encode(example) for example in examples]
        if subbatch_max_tokens is not None:
            res = subbatching.map(
                self._parse_encoded,
                examples,
                encoded,
                costs=self._get_lens(encoded),
                max_cost=subbatch_max_tokens,
                return_compressed=return_compressed,
                return_scores=return_scores,
                return_cats=return_cats,
                tau=tau
            )
        else:
            res = self._parse_encoded(
                examples,
                encoded,
                return_compressed=return_compressed,
                return_scores=return_scores,
                return_cats=return_cats,
                tau=tau
            )
            # fixing determinism bug
            res = list(res)
        self.train(training)
        if return_encoded:
            return res, encoded
        return res


text:
# coding=utf-8
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

import datetime
import os
import sys
import numpy as np
import cv2

from StreamManagerApi import StreamManagerApi
from StreamManagerApi import MxDataInput
from StreamManagerApi import StringVector
from StreamManagerApi import MxProtobufIn
from StreamManagerApi import InProtobufVector
import MxpiDataType_pb2 as MxpiDataType

if __name__ == '__main__':
    # init stream manager
    stream_manager_api = StreamManagerApi()
    ret = stream_manager_api.InitManager()
    if ret != 0:
        print("Failed to init Stream manager, ret=%s" % str(ret))
        exit()

    # create streams by pipeline config file
    with open("./mcnn.pipeline", 'rb') as f:
        pipelineStr = f.read()
    ret = stream_manager_api.CreateMultipleStreams(pipelineStr)
    if ret != 0:
        print("Failed to create Stream, ret=%s" % str(ret))
        exit()

    # Construct the input of the stream
    data_input = MxDataInput()

    dir_name = sys.argv[1]
    gt_name = sys.argv[2]

    file_list = os.listdir(dir_name)
    file_list.sort()
    mae = 0
    mse = 0
    start_time = datetime.datetime.now()
    for file_name in file_list:
        file_path = os.path.join(dir_name, file_name)
        gt_path = os.path.join(gt_name, file_name[:-3] + 'csv')
        if not (file_name.lower().endswith(".jpg") or file_name.lower().endswith(".jpeg") \
                or file_name.lower().endswith(".png")):
            continue

        empty_data = []
        stream_name = b'mcnn_opencv'
        in_plugin_id = 0
        input_key = 'appsrc0'

        img = cv2.imread(file_path, 0)
        img = img.astype(np.float32, copy=False)
        ht = img.shape[0]
        wd = img.shape[1]
        hang_left = 512 - int(ht / 2)
        hang_right = 1024 - hang_left - ht
        lie_left = 512 - int(wd / 2)
        lie_right = 1024 - lie_left - wd
        img = np.pad(img, ((hang_left, hang_right), (lie_left, lie_right)), 'constant')

        img = img.reshape((1, 1, 1024, 1024))
        tensor_list = MxpiDataType.MxpiTensorPackageList()
        tensor_pkg = tensor_list.tensorPackageVec.add()

        tensor_vec = tensor_pkg.tensorVec.add()
        tensor_vec.memType = 0
        tensor_vec.tensorShape.extend(img.shape)
        tensor_vec.tensorDataType = 0
        tensor_vec.dataStr = img.tobytes()
        tensor_vec.tensorDataSize = len(img)
        buf_type = b"MxTools.MxpiTensorPackageList"

        protobuf = MxProtobufIn()
        protobuf.key = input_key.encode("utf-8")
        protobuf.type = buf_type
        protobuf.protobuf = tensor_list.SerializeToString()
        protobuf_vec = InProtobufVector()
        protobuf_vec.push_back(protobuf)
        err_code = stream_manager_api.SendProtobuf(stream_name, in_plugin_id, protobuf_vec)
        if err_code != 0:
            print(
                "Failed to send data to stream, stream_name(%s), plugin_id(%s), element_name(%s), "
                "buf_type(%s), err_code(%s).", stream_name, in_plugin_id,
                input_key, buf_type, err_code)

        keys = [b"mxpi_tensorinfer0",]
        keyVec = StringVector()
        for key in keys:
            keyVec.push_back(key)
        infer_result = stream_manager_api.GetProtobuf(stream_name, 0, keyVec)
        if infer_result.size() == 0:
            print("infer_result is null")
            exit()

        if infer_result[0].errorCode != 0:
            print("GetProtobuf error. errorCode=%d" % (
                infer_result[0].errorCode))
            exit()

        TensorList = MxpiDataType.MxpiTensorPackageList()
        TensorList.ParseFromString(infer_result[0].messageBuf)
        data = np.frombuffer(TensorList.tensorPackageVec[0].tensorVec[0].dataStr, dtype=np.float32)

        den = np.loadtxt(open(gt_path, "rb"), delimiter=",", skiprows=0)
        den = den.astype(np.float32, copy=False)
        gt_count = np.sum(den)
        et_count = np.sum(data)
        mae += abs(gt_count - et_count)
        mse += ((gt_count - et_count) * (gt_count - et_count))
        print(file_path, "True value:", np.sum(den), "predictive value:", np.sum(data))

    mae = mae / 182
    mse = np.sqrt(mse / 182)
    end_time = datetime.datetime.now()
    print("*********************************************")
    print("Final accuracy of the project:")
    print('MAE:', mae, '  MSE:', mse)
    print("*********************************************")
    print("Overall project performance:")
    print(182 / (end_time - start_time).seconds, "images/seconds")

    # destroy streams
    stream_manager_api.DestroyAllStreams()

ext: py

max_stars_repo_name: mindspore-ai/models

max_issues_repo_name: mindspore-ai/models

max_forks_repo_name: mindspore-ai/models

reason: import numpy


text:
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
import argparse
import logging
import struct
import six

import pytest
import nose
import numpy as np

sys.path.append("../..")
from src.test_case import Predictor

FORMAT = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)

class TestModelInferenceTrtFp32(object):
    """
    TestModelInferenceTrtFp32
    Args:
    Returns:
    """

    def __init__(self,
                 data_path="Data/python-model-infer",
                 trt_dynamic_shape_info=None):
        """__init__
        """
        project_path = os.environ.get("project_path")
        self.model_root = os.path.join(project_path, data_path)
        self.trt_dynamic_shape_info = trt_dynamic_shape_info

    def check_data(self, result, expect, delta):
        """
        check result
        Args:
            result(list): list of result data
            expect(list): list of expect data
            delta(float): e.g. 0.001
        Returns:
            None
        """
        logger.info("current comparison delta is : {0}".format(delta))
        assert len(result) == pytest.approx(len(
            expect)), "output length not equal"
        for i in range(0, len(expect)):
            assert result[i] == pytest.approx(
                expect[i], rel=None, abs=delta), "output data not equal"

    def get_infer_results(self, model_path, data_path, min_subgraph_size=3):
        """
        get native and analysis infer results
        trt_fp32
        Args:
            model_path(string): parent path of __model__ file
            data_path(string): path of data.json
            min_subgraph_size(int): tensorrt subgraph size
        Returns:
            res(numpy array): analysis cf outputs
            exp(numpy array): native cfg outputs
        """
        AnalysisPredictor = Predictor(
            model_path,
            predictor_mode="Analysis",
            config_type="trt_fp32",
            min_subgraph_size=min_subgraph_size,
            trt_dynamic_shape_info=self.trt_dynamic_shape_info)

        res, ave_time = AnalysisPredictor.analysis_predict(
            data_path, repeats=10)
        logger.info(ave_time)

        NoIrPredictor = Predictor(
            model_path, predictor_mode="Analysis", config_type="gpu_no_ir")
        exp, ave_time = NoIrPredictor.analysis_predict(data_path)
        logger.info(ave_time)

        assert len(res) == pytest.approx(len(
            exp)), "num of output tensor not equal"

        return res, exp


text:
from collections import Counter
import random
import pickle as pkl
import datetime
import numpy as np

from playground.simulator import Simulator
from playground.state import State
from playground.selector import StreamWindowSelector, StreamSelector
from playground.strategy import Strategy
from playground.search import grid_search

class SSQ(State):

	def __init__(self,
				date,
				red_balls,
				blue_ball,
				sales=None,
				pool_size=None,
				price_count=None,
				price_money=None
				):
		self.date = date
		self.red_balls = red_balls
		self.blue_ball = blue_ball
		self.pool_size = pool_size
		self.sales = sales
		self.price_count = price_count
		self.price_money = price_money

	@classmethod
	def from_prediction(cls, date, red_balls, blue_ball):
		return cls(date=date, red_balls=red_balls, blue_ball=blue_ball)

	@classmethod
	def random(cls, date, seed=None):
		random.seed(seed)
		red_balls = random.sample(range(1, 34), k=6)
		blue_ball = random.choice(range(1,  17))
		return cls.from_prediction(date, red_balls, blue_ball)

	def check_result(self, other):
		if self.date != other.date:
			raise ValueError('Not from same date')

		n_match_red = self.number_of_matches(self.red_balls, other.red_balls)
		is_match_blue = (self.blue_ball == other.blue_ball)

		if n_match_red == 6 and is_match_blue:
			# 一等奖
			return self.price_money[1]
		elif n_match_red == 6 and not is_match_blue:
			# 二等奖
			return self.price_money[2]
		elif n_match_red == 5 and is_match_blue:
			# 三等奖
			return self.price_money[3]
		elif (n_match_red == 5 and not is_match_blue) or (n_match_red == 4 and is_match_blue):
			# 四等奖
			return self.price_money[4]
		elif (n_match_red == 4 and not is_match_blue) or (n_match_red == 3 and is_match_blue):
			# 五等奖
			return self.price_money[5]
		elif n_match_red in (0, 1, 2) and is_match_blue:
			# 六等奖
			return self.price_money[6]
		else:
			return 0

	def number_of_matches(self, a, b):
		return len(set(a) & set(b))

	def __repr__(self):
		return '<date: {} red: {} blue: {}>'.format(self.date, self.red_balls, self.blue_ball)

def load_history(data_path=None, start_date=None, end_date=None):
	""" Get the actual ticket result between start date and end date (inclusive),
		start_date and end_date can be datetime object of string of format 'yyyy-mm-dd'

		ex. load_history(start_date='2019-09-01', end_date='2019-10-01')  # 12 records
	"""
	start_date = datetime.datetime.strptime(start_date, '%Y-%m-%d').date() if isinstance(start_date, str) else start_date
	end_date = datetime.datetime.strptime(end_date, '%Y-%m-%d').date() if isinstance(end_date, str) else end_date

	data_path = data_path or './data/shuangseqiu.pkl'
	with open(data_path, 'rb') as f:
		history = pkl.load(f)

	history_states = [SSQ(**i) for i in history \
						if (start_date is None or start_date <= i['date']) and (end_date is None or end_date >= i['date'])]
	return history_states

#########################################
# Implementing some basic strategy ######
#########################################

COST_PER_TICKET = 2

class RandomSelectingStrategy(Strategy):

	def predict(self, env):
		cur_date = env['order_field']
		return SSQ.random(cur_date), COST_PER_TICKET

class FrequencyBasedStrategy(Strategy):

	def __init__(self,
				lookback_period=10,
				red_block=None):
		self.lookback_period = lookback_period
		if isinstance(red_block, int):
			if red_block == 6:
				red_block = [[0], [1], [2], [3], [4], [5]]
			elif red_block == 3:
				red_block = [[0, 1], [2, 3], [4, 5]]
			elif red_block == 2:
				red_block = [[0, 1, 2], [3, 4, 5]]
		self.red_block = red_block or [[0, 1, 2, 3, 4, 5]]

	def predict(self, env):
		cur_date = env['order_field']
		prev_states = env['state_info'][0]  # state_info is a tuple
		prev_states = prev_states[-self.lookback_period:]

		hist_blue_balls = [i.blue_ball for i in prev_states]
		pred_blue_ball = self.find_most_frequent(hist_blue_balls)[0][0]

		hist_red_balls = np.array([i.red_balls for i in prev_states])  # shape (lookback_period, 6)

		pred_red_balls = []
		for block in self.red_block:
			block_size = len(block)
			pred_block = self.find_most_frequent(hist_red_balls[:, block], n=block_size)
			pred_red_balls.extend([i[0] for i in pred_block])

		return SSQ.from_prediction(date=cur_date, red_balls=pred_red_balls, blue_ball=pred_blue_ball), COST_PER_TICKET

	@staticmethod
	def find_most_frequent(x, n=1):
		if hasattr(x, 'ndim') and x.ndim > 1:
			x = x.flatten()
		return Counter(x).most_common(n)

if __name__ == '__main__':
	states = load_history(start_date='2016-01-01')

	stream_selector = StreamWindowSelector(states, window_size=20, order_field='date')
	# strategy = RandomSelectingStrategy()
	strategy = FrequencyBasedStrategy(lookback_period=5, red_block=None)

	simulator = Simulator(stream_selector, strategy)
	simulator.simulate(steps=None)

	# params = {'lookback_period': np.arange(5, 100),
	# 		 'red_block': [None, 2, 3, 6]}

	# best_profit, best_param = grid_search(stream_selector, FrequencyBasedStrategy, params, steps=None)
	# print(best_profit)
	# print(best_param)


text:
# -*- coding: utf-8 -*-

import os
import json
import nltk
import gensim
import numpy as np
from gensim import corpora, models, similarities
import pickle

os.chdir("D:\semicolon\Deep Learning\chatbot");
model = gensim.models.Word2Vec.load('word2vec.bin');
path2="corpus";
file=open(path2+'/conversation.json');
data = json.load(file)
cor=data["conversations"];

x=[]
y=[]

path2="corpus";

for i in range(len(cor)):
    for j in range(len(cor[i])):
        if j<len(cor[i])-1:
            x.append(cor[i][j]);
            y.append(cor[i][j+1]);

tok_x=[]
tok_y=[]
for i in range(len(x)):
    tok_x.append(nltk.word_tokenize(x[i].lower()))
    tok_y.append(nltk.word_tokenize(y[i].lower()))

sentend=np.ones((300L,),dtype=np.float32)

vec_x=[]
for sent in tok_x:
    sentvec = [model[w] for w in sent if w in model.vocab]
    vec_x.append(sentvec)

vec_y=[]
for sent in tok_y:
    sentvec = [model[w] for w in sent if w in model.vocab]
    vec_y.append(sentvec)

for tok_sent in vec_x:
    tok_sent[14:]=[]
    tok_sent.append(sentend)

for tok_sent in vec_x:
    if len(tok_sent)<15:
        for i in range(15-len(tok_sent)):
            tok_sent.append(sentend)

for tok_sent in vec_y:
    tok_sent[14:]=[]
    tok_sent.append(sentend)

for tok_sent in vec_y:
    if len(tok_sent)<15:
        for i in range(15-len(tok_sent)):
            tok_sent.append(sentend)

with open('conversation.pickle','w') as f:
    pickle.dump([vec_x,vec_y],f)


text:
import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from functools import reduce
import os, sys

def compute_std_one_interaction(df, offset, gene_A, gene_B):
    """
    df: dataframe with data (sample x gene).
    offset: matrix of sample offset.
    gene_A, gene_B: two genes to screen for
    """
    inter_std = np.multiply(
        np.multiply(df[gene_A], df[gene_B]),
        offset
    )
    return [gene_A, gene_B, np.std(inter_std)]

def _compute_all_interactions_std_one_gene(df, offset, gene_A):
    return [
        compute_std_one_interaction(df, offset, gene_A, gene_B)
        for gene_B in df.columns if gene_A <= gene_B
    ]

def compute_all_interactions_std(df, offset):
    return Parallel(n_jobs=30, verbose=1)(
        delayed(_compute_all_interactions_std_one_gene)(df, offset, gene_A)
        for gene_A in df.columns
    )

def read_interaction_weights(
    output_folder,
    iter_idx,
    file_radical,
    kernel_name='laplacian'
):
    """
    Read the interaction terms of two folders: one with Gaussian and one with Laplacian kernel.
    Proceed to correct the weights to harmonize the norms.

    - output_folder: Output folder of the "feature_analysis" experiment.
    - iter_idx: Iteration ID considered.
    - file_radical: Name of the file to consider.
    """
    # Read interaction terms
    print('START READING INTERACTIONS %s'%(file_radical), flush=True)
    PV_inter_features = _read_interaction_terms(
        iter_idx=iter_idx,
        file_radical=file_radical,
        output_folder=output_folder,
        kernel_name=kernel_name
    )
    PV_inter_features = _process_interaction_terms(PV_inter_features)

    # Aggregate scaling coefficients (to re-use later in permutation testing)
    source_interaction_scaling_coef = PV_inter_features[['kernel', 'factor', 'multiplying_factor', 'is_squared']]
    source_interaction_scaling_coef = source_interaction_scaling_coef.drop_duplicates()
    source_interaction_scaling_coef = source_interaction_scaling_coef.set_index('kernel')
    source_interaction_scaling_coef = source_interaction_scaling_coef.loc[kernel_name]

    # Restrict to Laplacian and return both the scaling coefficients of Laplacian kernel
    # and the corrected values in a gene x PV matrix
    return source_interaction_scaling_coef, _restrict_laplacian_kernel(PV_inter_features, kernel_name)

def _read_interaction_terms(iter_idx, file_radical, output_folder, kernel_name='laplacian'):
    """
    Read one interaction files and correct between Gaussian and Laplacian terms.
    """

    # Read the different interaction terms
    feature_df = pd.concat({
        k: pd.read_pickle(
            '%s/iter_%s_nu_%s/%s.gz'%(output_folder, iter_idx, k, file_radical),
            compression='gzip'
        ) for k in [kernel_name, 'gaussian']
    }).T
    feature_df = feature_df.reset_index().melt(id_vars='index')
    feature_df.columns = ['interaction', 'kernel', 'factor', 'weight']
    feature_df['square_weight'] = np.square(feature_df['weight'])
    feature_df['is_squared'] = feature_df['interaction'].str.contains(r'(\^2)')

    # # Compute Laplacian reweighting
    norm_df = feature_df.groupby(['kernel', 'factor', 'is_squared']).agg('sum')
    norm_df = norm_df.pivot_table(columns='kernel', index=['factor', 'is_squared'])['square_weight']
    norm_df['multiplying_factor'] = norm_df['gaussian'] / norm_df[kernel_name]
    norm_df['multiplying_factor'] = np.sqrt(norm_df['multiplying_factor'])

    # Aggregate with all weights
    feature_df = feature_df.merge(norm_df, on=['factor', 'is_squared'])
    del feature_df['gaussian'], feature_df[kernel_name]#, feature_df['is_squared']
    feature_df.loc[feature_df['kernel'] == 'gaussian', 'multiplying_factor'] = 1
    feature_df['corrected_weight'] = feature_df['weight'] * feature_df['multiplying_factor']
    feature_df['square_corrected_weight'] = np.square(feature_df['corrected_weight'])

    return feature_df

def _process_interaction_terms(PV_inter_features):
    # Process gene name
    PV_inter_features[['gene_A', 'gene_B']] = PV_inter_features['interaction'].str.split('*', expand=True)
    is_square_gene = pd.isnull(PV_inter_features['gene_B'])
    PV_inter_features.loc[is_square_gene, 'gene_B'] = PV_inter_features.loc[is_square_gene, 'gene_A']
    for x in ['gene_A', 'gene_B']:
        PV_inter_features[x] = PV_inter_features[x].str.replace('(\^[1-9])$', '')
    del PV_inter_features['interaction']

    return PV_inter_features

def _restrict_laplacian_kernel(PV_inter_features, kernel_name='laplacian'):
    # Restrict to Laplacian features
    laplacian_PV_inter_features = PV_inter_features.loc[PV_inter_features['kernel'] == kernel_name]
    laplacian_PV_inter_features = laplacian_PV_inter_features[['factor', 'gene_A', 'gene_B', 'corrected_weight']]

    return laplacian_PV_inter_features.pivot_table(
        values='corrected_weight',
        columns='factor',
        index=['gene_A', 'gene_B']
    )

def _compute_kernel_param(data_source, sobolev_alignment_clf):
    clf = sobolev_alignment_clf['gaussian']
    sigma = clf.krr_params[data_source]['kernel_params']['sigma']
    return 1 / (2 * sigma**2)

def read_and_process_sample_permutated_interactions(
    output_folder,
    iter_idx,
    file_radical,
    interaction_scaling_coef,
    kernel_name='laplacian'
):
    """
    Read, process and correct for Laplacian difference for one sample-permuted file.

    - output_folder: folder of one "feature_analysis" experiment.

    - iter_idx: index of the iteration to study.

    - file_radical: file to look for inside the iteration subfolder.

    - interaction_scaling_coef: scaling coefficients DataFrame for Laplacian kernel (given as
    first output of read_interaction_weights)
    """
    # Read interactions file
    permuted_interactions_df = pd.read_pickle(
        '%s/GSEA_null/iter_%s_nu_%s/%s.gz'%(output_folder, iter_idx, kernel_name, file_radical),
        compression='gzip'
    )

    # Processing
    permuted_interactions_df = permuted_interactions_df.reset_index().melt(id_vars='index')
    permuted_interactions_df.columns = ['factor', 'interaction', 'weight']
    permuted_interactions_df['square_weight'] = np.square(permuted_interactions_df['weight'])
    permuted_interactions_df['is_squared'] = permuted_interactions_df['interaction'].str.contains(r'(\^2)')

    # Merge with scaling for interactions
    permuted_interactions_df = permuted_interactions_df.merge(
        interaction_scaling_coef,
        on=['factor', 'is_squared']
    )
    permuted_interactions_df['corrected_weight'] = permuted_interactions_df['weight'] * permuted_interactions_df['multiplying_factor']
    permuted_interactions_df['square_corrected_weight'] = np.square(permuted_interactions_df['corrected_weight'])

    # Format interactions
    permuted_interactions_df = _process_interaction_terms(permuted_interactions_df)
    permuted_interactions_df = permuted_interactions_df[['factor', 'gene_A', 'gene_B', 'corrected_weight']]

    return permuted_interactions_df.pivot_table(
        values='corrected_weight',
        columns='factor',
        index=['gene_A', 'gene_B']
    )


text:
"""
File: run_text_classification.py
Author: ANONYMIZED
Email: ANONYMIZED
Github: ANONYMIZED
Description: Run text classification experiments on TextGCN's datasets
"""

import csv
import itertools as it
import logging

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import numpy as np

import scipy.sparse as sp
import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
from joblib import Memory
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm, trange
from transformers import (WEIGHTS_NAME, AdamW, AutoTokenizer, BertConfig,
                          BertForSequenceClassification, BertModel,
                          BertTokenizer, DistilBertConfig,
                          DistilBertForSequenceClassification, DistilBertModel,
                          DistilBertTokenizer, get_linear_schedule_with_warmup)

from transformers import DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,\
    BERT_PRETRAINED_MODEL_ARCHIVE_LIST

from sklearn.metrics import f1_score
from sklearn.feature_extraction.text import TfidfTransformer

from tokenization import build_tokenizer_for_word_embeddings
from data import load_data, load_word_vectors, shuffle_augment
from models import MLP, collate_for_mlp

try:
    import wandb
    WANDB = True
except ImportError:
    print("WandB not installed, to track experiments: pip install wandb")
    WANDB = False

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)

logger = logging.getLogger(__name__)
USE_CUDA = torch.cuda.is_available()
CACHE_DIR = 'cache/textclf'
MEMORY = Memory(CACHE_DIR, verbose=2)

VALID_DATASETS = [ '20ng', 'R8', 'R52', 'ohsumed', 'mr'] + ['TREC', 'wiki']

# ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) for conf in (BertConfig, XLNetConfig, XLMConfig,
#                                                                                RobertaConfig, DistilBertConfig)), ())

ALL_MODELS = BERT_PRETRAINED_MODEL_ARCHIVE_LIST + DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST
MODEL_CLASSES = {
    'bert': (BertConfig, BertForSequenceClassification, BertModel),
    # 'xlnet': (XLNetConfig, XLNetForSequenceClassification, XLNetModel),
    # 'xlm': (XLMConfig, XLMForSequenceClassification, XLMModel),
    # 'roberta': (RobertaConfig, RobertaForSequenceClassification, RobertaModel),
    'distilbert': (DistilBertConfig, DistilBertForSequenceClassification, DistilBertModel)
}

def inverse_document_frequency(encoded_docs, vocab_size):
    """ Returns IDF scores in shape [vocab_size] """
    num_docs = len(encoded_docs)
    counts = sp.dok_matrix((num_docs, vocab_size))
    for i, doc in tqdm(enumerate(encoded_docs), desc="Computing IDF"):
        for j in doc:
            counts[i,j] += 1

    tfidf = TfidfTransformer(use_idf=True, smooth_idf=True)

    tfidf.fit(counts)

    return torch.FloatTensor(tfidf.idf_)

def pad(seqs, with_token=0, to_length=None):
    if to_length is None:
        to_length = max(len(seq) for seq in seqs)
    return [seq + (to_length - len(seq)) * [with_token] for seq in seqs]

def get_collate_for_transformer(pad_token_id):
    """ Closure to include padding in collate function """
    def _collate_for_transformer(examples):
        docs, labels = list(zip(*examples))
        input_ids = torch.tensor(pad(docs, with_token=pad_token_id))
        attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        attention_mask[input_ids == pad_token_id] = 0
        labels = torch.tensor(labels)
        token_type_ids = torch.zeros_like(input_ids)
        return input_ids, attention_mask, token_type_ids, labels
    return _collate_for_transformer

def train(args,  train_data, model, tokenizer):
    if args.model_type == 'mlp':
        # if args.model_name_or_path is not None:
        #     # Embedding case, use optimized merging from tokenizer lib
        #     collate_fn = collate_encoding_for_mlp
        # else:
        #     # Manual collate with offsets
        collate_fn = collate_for_mlp
    else:
        collate_fn = get_collate_for_transformer(tokenizer.pad_token_id)
    train_loader = torch.utils.data.DataLoader(train_data,
                                               collate_fn=collate_fn,
                                               shuffle=True,
                                               batch_size=args.batch_size,
                                               num_workers=args.num_workers,
                                               pin_memory=('cuda' in str(args.device)))

    # len(train_loader) no of batches
    t_total = len(train_loader) // args.gradient_accumulation_steps * args.epochs

    optimizer = AdamW(model.parameters(), lr=args.learning_rate, eps=args.adam_epsilon)
    # scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
    writer = SummaryWriter()

    if args.ignore_position_ids:
        print("Setting position ids to zero and ignoring grad")
        if args.model_type == 'bert':
            model.bert.embeddings.position_embeddings.weight.requires_grad = False
            model.bert.embeddings.position_embeddings.weight.zero_()
        else:
            raise NotImplementedError("Ignore position ids only implemented for BERT")

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_data))
    logger.info("  Num Epochs = %d", args.epochs)
    logger.info("  Batch size  = %d", args.batch_size)
    logger.info("  Total train batch size (w. accumulation) = %d",
                   args.batch_size * args.gradient_accumulation_steps)
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(args.epochs, desc="Epoch")
    for epoch in train_iterator:
        epoch_iterator = tqdm(train_loader, desc="Iteration")
        for step, batch in enumerate(epoch_iterator):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            if args.model_type == 'mlp':
                # Batch: torch.tensor(flat_docs), torch.tensor(offsets), torch.tensor(labels)
                outputs = model(batch[0], batch[1], batch[2])
            else:
                # Batch : input_ids, attention_mask, token_type_ids, labels
                inputs = {'input_ids':      batch[0],
                          'attention_mask': batch[1],
                          'labels':         batch[3]}
                if args.model_type != 'distilbert':
                    inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
                if args.ignore_position_ids:
                    inputs['position_ids'] = torch.zeros(
                        inputs['input_ids'].shape[0],  # bsz
                        inputs['input_ids'].shape[1],  # len
                        device=inputs['input_ids'].device,
                        dtype=torch.long
                    )
                outputs = model(**inputs)
            loss = outputs[0]
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
            loss.backward()
            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
                optimizer.step()
                scheduler.step()
                model.zero_grad()
                global_step += 1
                if WANDB:
                    wandb.log({'epoch': epoch,
                               'lr': scheduler.get_last_lr()[0],
                               'loss': loss})

            if args.logging_steps > 0 and global_step % args.logging_steps == 0:
                # if args.evaluate_during_training:
                #     results = evaluate(args, dev_data, model, tokenizer)
                #     for key, value in results.items():
                #         tb_writer.add_scalar('eval_{}'.format(key), value, global_step)

                avg_loss = (tr_loss - logging_loss)/ args.logging_steps
                writer.add_scalar('lr', scheduler.get_last_lr()[0], global_step)
                writer.add_scalar('loss', avg_loss, global_step)
                logging_loss = tr_loss

    writer.close()
    return global_step, tr_loss / global_step

def evaluate(args, dev_or_test_data, model, tokenizer):
    if args.model_type == 'mlp':
        collate_fn = collate_for_mlp
    else:
        collate_fn = get_collate_for_transformer(tokenizer.pad_token_id)
    data_loader = torch.utils.data.DataLoader(dev_or_test_data,
                                              collate_fn=collate_fn,
                                              num_workers=args.num_workers,
                                              batch_size=args.test_batch_size,
                                              pin_memory=('cuda' in str(args.device)),
                                              shuffle=False)
    all_logits = []
    all_targets = []
    nb_eval_steps, eval_loss = 0, 0.0
    for batch in tqdm(data_loader, desc="Evaluating"):
        model.eval()
        batch = tuple(t.to(args.device) for t in batch)
        with torch.no_grad():
            if args.model_type == 'mlp':
                # batch consist of (flat_inputs, lenghts, labels)
                outputs = model(batch[0], batch[1], batch[2])
                all_targets.append(batch[2].detach().cpu())
            else:
                inputs = {'input_ids':      batch[0],
                          'attention_mask': batch[1],
                          'labels':         batch[3]}
                if args.model_type != 'distilbert':
                    inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
                outputs = model(**inputs)
                all_targets.append(inputs['labels'].detach().cpu())
        nb_eval_steps += 1
        # outputs [:2] should hold loss, logits
        loss, logits = outputs[:2]
        eval_loss += loss.mean().item()
        all_logits.append(logits.detach().cpu())

    logits = torch.cat(all_logits).numpy()
    targets = torch.cat(all_targets).numpy()
    eval_loss /= nb_eval_steps
    preds = np.argmax(logits, axis=1)
    acc = (preds == targets).sum() / targets.size

    f1_micro = f1_score(targets, preds, average='micro')
    f1_macro = f1_score(targets, preds, average='macro', zero_division=1)

    if WANDB:
        wandb.log({"test/acc": acc, "test/loss": eval_loss,
                   "test/f1_micro": f1_micro,
                   "test/f1_macro": f1_macro})
    return acc, eval_loss

def run_xy_model(args):
    print("Loading data...")

    if args.model_type == "mlp" and args.model_name_or_path is not None:
        print("Assuming to use word embeddings as both model_type=mlp and model_name_or_path are given")
        print("Using word embeddings -> forcing wordlevel tokenizer")
        vocab, embedding = load_word_vectors(args.model_name_or_path, unk_token="[UNK]")
        tokenizer = build_tokenizer_for_word_embeddings(vocab)
    else:
        tokenizer_name = args.tokenizer_name if args.tokenizer_name else args.model_name_or_path
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        embedding = None
    print("Using tokenizer:", tokenizer)

    do_truncate = not (args.stats_and_exit or args.model_type == 'mlp')
    if args.stats_and_exit:
        # We only compute dataset stats including length, so NOT truncate
        max_length = None
    elif args.model_type == 'mlp':
        max_length = None
    else:
        max_length = 512 # should hold for all used transformer models?

    enc_docs, enc_labels, train_mask, test_mask, label2index = load_data(args.dataset,
                                                                         tokenizer,
                                                                         max_length=max_length,
                                                                         construct_textgraph=False,
                                                                         n_jobs=args.num_workers)
    print("Done")

    lens = np.array([len(doc) for doc in enc_docs])
    print("Min/max document length:", (lens.min(), lens.max()))
    print("Mean document length: {:.4f} ({:.4f})".format(lens.mean(), lens.std()))
    assert len(enc_docs) == len(enc_labels) == train_mask.size(0) == test_mask.size(0)
    enc_docs_arr, enc_labels_arr = np.array(enc_docs, dtype='object'), np.array(enc_labels)

    train_docs = enc_docs_arr[train_mask]
    train_labels = enc_labels_arr[train_mask]

    if args.shuffle_augment:
        factor = float(args.shuffle_augment)

        # Generate new permuted documents
        new_docs, new_labels = shuffle_augment(list(train_docs),
                                               list(train_labels),
                                               factor=factor,
                                               random_seed=args.seed)

        # Convert to numpy
        new_docs = np.array(new_docs, dtype='object')
        new_labels = np.array(new_labels)

        # Augment the training data
        train_docs = np.concatenate([train_docs, new_docs])
        train_labels = np.concatenate([train_labels, new_labels])

    train_data = list(zip(train_docs, train_labels))

    test_data = list(zip(enc_docs_arr[test_mask], enc_labels_arr[test_mask]))

    print("N", len(enc_docs))
    print("N train", len(train_data))
    print("N test", len(test_data))
    print("N classes", len(label2index))

    if args.stats_and_exit:
        print("Warning: length stats depend on tokenizer and max_length of model, chose MLP to avoid trimming before computing stats.")
        exit(0)

    if args.model_type != 'mlp':
        config_class, model_class, __ = MODEL_CLASSES[args.model_type]
        print("Loading", args.model_type)
        print("Loading config")
        config = config_class.from_pretrained(args.model_name_or_path,
                                              num_labels=len(label2index),
                                              cache_dir=CACHE_DIR)

        print(config)
        print("Loading model")
        model = model_class.from_pretrained(args.model_name_or_path,
                                            from_tf=bool('.ckpt' in args.model_name_or_path),
                                            config=config,
                                            cache_dir=CACHE_DIR)
        # This is a ForSequenceClassification Model
    else:
        print("Initializing MLP")

        if embedding is not None:
            # Vocab size given by embedding
            vocab_size = None
        else:
            vocab_size = tokenizer.vocab_size

        if args.bow_aggregation == 'tfidf':
            print("Using IDF")
            idf = inverse_document_frequency(enc_docs_arr[train_mask], tokenizer.vocab_size).to(args.device)
        else:
            idf = None

        model = MLP(vocab_size, len(label2index),
                    num_hidden_layers=args.mlp_num_layers,
                    hidden_size=args.mlp_hidden_size,
                    embedding_dropout=args.mlp_embedding_dropout,
                    dropout=args.mlp_dropout,
                    mode=args.bow_aggregation,
                    pretrained_embedding=embedding,
                    idf=idf,
                    freeze=args.freeze_embedding)

    model.to(args.device)

    if WANDB:
        wandb.watch(model, log_freq=args.logging_steps)

    train(args, train_data, model, tokenizer)
    acc, eval_loss = evaluate(args, test_data, model, tokenizer)
    print(f"[{args.dataset}] Test accuracy: {acc:.4f}, Eval loss: {eval_loss}")
    return acc

def main():
    import argparse
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('dataset', choices=VALID_DATASETS)
    parser.add_argument("--model_type", default=None, type=str, required=True,
                        help="Model type: either 'mlp' or 'distilbert'",
                        choices=["mlp", "distilbert", "bert"])
    parser.add_argument("--model_name_or_path", default=None, type=str,
                        help="Optional path to word embedding with model type 'mlp' OR huggingface shortcut name such as distilbert-base-uncased for model type 'distilbert'")
    parser.add_argument("--results_file", default=None,
                        help="Store results to this results file")

    ## Training config
    parser.add_argument("--epochs", type=int, default=20,
                        help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=16,
                        help="Batch size for training")
    parser.add_argument("--test_batch_size", type=int, default=None,
                        help="Batch size for testing (defaults to train batch size)")
    parser.add_argument('--logging_steps', type=int, default=50,
                        help="Log every X updates steps.")

    parser.add_argument('--unfreeze_embedding', dest="freeze_embedding", default=True,
            action='store_false', help="Allow updating pretrained embeddings")

    ## Training Hyperparameters
    parser.add_argument("--learning_rate", default=5e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument("--weight_decay", default=0.0, type=float,
                        help="Weight deay if we apply some.")
    parser.add_argument("--warmup_steps", default=0, type=int,
                        help="Linear warmup over warmup_steps.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
                        help="Epsilon for Adam optimizer.")

    ## Other parameters
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument("--num_workers", default=4, type=int,
                        help="Number of workers")

    parser.add_argument("--stats_and_exit", default=False,
                        action='store_true',
                        help="Print dataset stats and exit.")

    # MLP Params
    parser.add_argument("--mlp_num_layers", default=1, type=int, help="Number of hidden layers within MLP")
    parser.add_argument("--mlp_hidden_size", default=1024, type=int, help="Hidden dimension for MLP")
    parser.add_argument("--bow_aggregation", default="mean", choices=["mean", "sum", "tfidf"],
            help="Aggregation for bag-of-words models (such as MLP)")
    parser.add_argument("--mlp_embedding_dropout", default=0.5, type=float, help="Dropout for embedding / first hidden layer ")
    parser.add_argument("--mlp_dropout", default=0.5, type=float, help="Dropout for all subsequent layers")

    parser.add_argument("--comment", help="Some comment for the experiment")
    parser.add_argument("--ignore_position_ids",
                        help="Use all zeros to pos ids",
                        default=False, action='store_true')
    parser.add_argument("--seed", default=None,
                        help="Random seed for shuffle augment")
    parser.add_argument("--shuffle_augment", type=float,
                        default=0, help="Factor for shuffle data augmentation")
    ##########################
    args = parser.parse_args()

    if args.model_type in ['mlp', 'textgcn']:
        assert args.tokenizer_name or args.model_name_or_path, "Please supply tokenizer for MLP via --tokenizer_name or provide an embedding via --model_name_or_path"
    else:
        assert args.model_name_or_path, f"Please supply --model_name_or_path for {args.model_type}"

    args.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    args.test_batch_size = args.batch_size if args.test_batch_size is None else args.test_batch_size

    if WANDB:
        wandb.init(project="text-clf")
        wandb.config.update(args)

    acc = {
        'mlp': run_xy_model,
        'bert': run_xy_model,
        'distilbert': run_xy_model,
        'roberta': run_xy_model,
        'xlnet': run_xy_model
    }[args.model_type](args)
    if args.results_file:
        with open(args.results_file, 'a', newline='') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow([args.model_type,args.dataset,acc])

if __name__ == '__main__':
    main()

max_stars_repo_name: lgalke/text-clf-baselines

max_issues_repo_name: lgalke/text-clf-baselines

max_forks_repo_name: lgalke/text-clf-baselines

reason: import numpy,import scipy


text:
"""
    Set of functions and classes that help dealing with creating, modifying and saving primitives, such as parameters or the utility functions.
"""

import numpy as np
import logging
import json

class Parameters:

    """
        Class whose objects store the parameters used in the model.
    """

    def __init__(self, overwrite={}):
        """
            We can either initialize the object using a pre-defined set of parameters, or pass in a user-defined set
             of parameters.
            :param model_input: Dict or None
        """

        # Points in the Model
        self.num_l  = 101     # Number of points of evaluation
        self.num_v  = 200     # Number of points in the grid for V
        self.num_x  = 15      # Number of points of support for worker productivity
        self.num_np = 5       # Number of non-permanent levels
        self.num_z  = 7       # Number of points for match productivity
        self.num_s  = 50      # Number of points of support for piece rate contract

        # Time periods in the Model
        self.dt     = 0.25    # Time as a Fraction of Year

        # Utility Function Parameters
        self.u_rho = 1.5       # Risk aversion coefficient
        self.u_a   = 1.0
        self.u_b   = 1.0

        # Search Environment
        self.z_0      = 4          # Slice of value function of firms (index starts at 1)
        self.s_job    = 1.0        # Relative Efficiency of Search on the Job
        self.alpha    = 0.1        # Parameter for probability of finding a job
        self.sigma    = 1.0        # Parameter for probability of finding a job
        self.kappa    = 1.0        # Vacancy cost parameter

        # effort function that control separation
        self.efcost_sep = 0.005 * self.dt
        self.efcost_ce  = 0.3

        # Productivity shocks
        self.x_corr = 0.95  # Correlation in worker productivity
        self.z_corr = 0.95  # Correlation in match productivity

        # Productivity Function Parameters
        self.prod_var_x  = 1.0           # Variance of X (permanent)
        self.prod_var_x2 = 1.0           # Variance of X (non-permanent)
        self.prod_var_z  = 1.0           # Variance of Z
        self.prod_z      = 0.5           # Production function parameter
        self.prod_rho    = 1.0           # Production function parameter
        self.prod_mu     = 0.2           # Worker contribution
        self.prod_px     = 1.0           # Worker power (non linear in type)
        self.prod_py     = 1.0           # Firm power (nonlinear in type)
        self.prod_a      = 16 * self.dt  # Factor for output function
        self.prod_err_w  = 0.0           # Measurement error on wages
        self.prod_err_y  = 0.0           # Measurement error on wages

        # Discounting Rates
        self.beta     = 1 - (1 - 0.95) * self.dt  # Impatience
        self.int_rate = 1 / self.beta - 1         # Period interest rate

        # Unemployment Parameters
        self.u_bf_m = 0.05       # Intercept of benefit function for unemployed(x)
        self.u_bf_c = 0.5        # Slope of benefit function for unemployed(x) not used

        # Unemployment Parameters w_net = tau * w ^ lambda
        self.tax_lambda = 1.0       # curvature of the tax system
        self.tax_tau    = 1.0       # proportion of take home
        self.tax_expost_lambda = 1.0  # this is for counterfactuals, allows to only apply taxes expost
        self.tax_expost_tau = 1.0     # this is for counterfactuals, allows to only apply taxes expost

        # Computational Parameters
        self.chain            = 1         # Chain id when running in parallel
        self.max_iter         = 5000
        self.max_iter_fb      = 5000
        self.verbose          = 5
        self.iter_display     = 25
        self.tol_simple_model = 1e-9
        self.tol_full_model   = 1e-8
        self.eq_relax_power   = 0.4       #  we relax the equilibrium constrain using an update rule based
        self.eq_relax_margin  = 500       #  on mumber of iterations
        self.eq_weighting_at0 = 0.01      # fitting J function with weight around 0

        # simulation parameters
        self.sim_ni      = 20000  # number of workers
        self.sim_nt      = 30     # time periods on top of nt_burn
        self.sim_nt_burn = 10     # periods to discard at begining
        self.sim_nh      = 200    # length of the firm history
        self.sim_nrep    = 20     # number of replication samples
        self.sim_net_earnings = False # whether to use net or gross earnings in the simulation

        for key, val in overwrite.items():
            if key in self.__dict__.keys():
                self.__dict__[key] = val
            else:
                logging.warning("this key does not exists:" + key)

    @staticmethod
    def load(filename) -> 'Parameters':
        with open(filename, "r") as infile:
            pdict =  json.load(infile)
        p = Parameters(pdict)
        return p

    def save(self,filename, append_dict = {}):
        temp_dict = self.__dict__.copy()
        temp_dict.update(append_dict)
        with open(filename, "w") as fp:
            json.dump(temp_dict, fp)

    def __getstate__(self):
        """ defines how the model is pickled """
        return self.__dict__     # get attribute dictionary

    def __setstate__(self, dict):
        """ defines how the model is unpickled """
        self.__dict__ = dict     # make dict our attribute dictionary

    def to_dict(self):
        return self.__dict__.copy()

    def get_x_components(self):
        """ give the values of x0 and x1 for each of the values of x """
        num_x0 = int(self.num_x / self.num_np)

        # the fixed heterogeneity component
        x0 = np.arange(num_x0)
        xt = np.arange(self.num_np)
        x0 = np.kron(x0,np.ones(self.num_np)) # permanent is slow moving
        xt = np.kron(np.ones(num_x0),xt) # permanent is slow moving
        return x0,xt

class Preferences:

    """
        Class whose methods represent the preferences and their derivatives, taking a Parameters object as input.
    """

    def __init__(self, input_param=None):
        """
            Initialize with a parameter object.
            :param input_param: Input parameter object, can be None
        """
        if input_param is None:
            self.p = Parameters()
        else:
            self.p = input_param

    def utility(self, wage):
        """
            Computes the utility function at a particular wage.
            :param wage: Argument of the function.
            :return: Output of the function.
        """
        aa = self.p.u_a * np.power(self.p.tax_tau, 1 - self.p.u_rho)
        return np.divide(aa * np.power( wage, self.p.tax_lambda * (1.0 - self.p.u_rho)) - self.p.u_b,
                         1 - self.p.u_rho)
        # return np.divide(self.p.u_a * np.power(wage, 1 - self.p.u_rho) - self.p.u_b,
        #                  1 - self.p.u_rho)

    def utility_gross(self, wage):
        """
            Computes the utility function at a particular wage, not applying the tax function
            :param wage: Argument of the function.
            :return: Output of the function.
        """
        return np.divide(self.p.u_a * np.power(wage, 1 - self.p.u_rho) - self.p.u_b,
                         1 - self.p.u_rho)

    def inv_utility(self, value):
        """
            Computes the inverse utility function at a particular value.
            :param value: Argument of the function.
            :return: Output of the function.
        """
        aa = self.p.u_a * np.power(self.p.tax_tau, 1.0 - self.p.u_rho)
        return np.power(np.divide((1.0 - self.p.u_rho) * value + self.p.u_b, aa),
                        (np.divide(1.0, self.p.tax_lambda * (1.0 - self.p.u_rho))))

    def utility_1d(self, wage):
        """
            Computes the first derivative of the utility function at a particular wage.
            :param wage: Argument of the function.
            :return: Output of the function.
        """
        #return self.p.u_a * np.power(wage, - self.p.u_rho)
        aa = self.p.u_a * np.power(self.p.tax_tau, 1.0 - self.p.u_rho)
        return aa * self.p.tax_lambda * np.power(wage, self.p.tax_lambda * ( 1.0 - self.p.u_rho) - 1.0)

    def inv_utility_1d(self, value):
        """
            Computes the first derivative of the inverse utility function at a particular value.
            :param value: Argument of the function.
            :return: Output of the function.
        """
        aa = self.p.u_a * np.power(self.p.tax_tau, 1 - self.p.u_rho)
        pow_arg = ( (1 - self.p.u_rho) * value + self.p.u_b   ) / aa
        return np.power( pow_arg, 1.0/( self.p.tax_lambda * (1 - self.p.u_rho) ) - 1.0) / ( self.p.tax_lambda * aa )

    def effort_cost(self, q_value):
        """
            Computes the effort cost function given the level of effort 'q'.
            :param q_value: Argument of the function.
            :return: Output of the function.
        """
        gam = self.p.efcost_ce ** -1
        return self.p.efcost_sep * q_value + \
            np.divide(self.p.efcost_sep,  self.p.efcost_ce - 1) - \
            np.divide(self.p.efcost_sep, 1 - gam) * (q_value ** (1 - gam))

    def inv_effort_cost_1d(self, V):
        """
            Returns the quit probability by computing the inverse of the first derivative of the
            effort cost function.
            :param eff: Argument of the function.
            :return: Output of the function.
        """
        #return np.power(1 + np.divide(np.maximum(eff, 0), self.p.efcost_sep), -self.p.efcost_ce)
        return np.power(1 + np.divide(np.maximum(-V, 0), self.p.efcost_sep), -self.p.efcost_ce)

    def log_consumption_eq(self, V):
        """
            Returns the log wage/consumption equivalent associated with a present value of the worker.
        """
        return(np.log(self.inv_utility( (1-self.p.beta) * V )))

    def log_profit_eq(self,J):
        """
            Returns the log profit equivalent associated with the firm present value
        """
        return( np.log( (1-self.p.beta) * J))

    def consumption_eq(self, V):
        """
            Returns the log wage/consumption equivalent associated with a present value of the worker.
        """
        return((self.inv_utility( (1-self.p.beta) * V )))

    def profit_eq(self,J):
        """
            Returns the log profit equivalent associated with the firm present value
        """
        return(( (1-self.p.beta) * J))


text:
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import pickle

customers = pd.read_csv("Ecommerce Customers")

y = customers['Yearly Amount Spent']
X = customers[['Avg. Session Length', 'Time on App','Time on Website', 'Length of Membership']]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=101)

lm = LinearRegression()
lm.fit(X_train,y_train)

filename = "model.pkl"
pickle.dump(lm, open(filename, 'wb'))