text:
import argparse
import glob
import json
import logging

import os
import random
import math
import copy
import numpy as np
import torch
# from torch.functional import norm
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
from nni.compression.pytorch.utils import get_module_by_name
from emmental import MaskedBertConfig, MaskedBertForSequenceClassification
from transformers import (
    WEIGHTS_NAME,
    AdamW,
    BertConfig,
    BertForSequenceClassification,
    BertTokenizer,
    get_linear_schedule_with_warmup,
)
from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes as output_modes
from transformers import glue_processors as processors
from emmental.modules.masked_nn import MaskedLinear

from nni.compression.pytorch import ModelSpeedup, apply_compression_results
from nni.algorithms.compression.pytorch.pruning import L1FilterPruner, LevelPruner
from nni.algorithms.compression.pytorch.pruning.weight_masker import WeightMasker
from nni.algorithms.compression.pytorch.pruning.dependency_aware_pruner import DependencyAwarePruner
from shape_hook import ShapeHook

head_pruner_cfg = torch.load('head_prune_cfg')
task_name = "qqp"
model_name_or_path = '../training/result/qqp_partial/coarse_0.3/checkpoint-220000/'
data_dir = './QQP'
max_seq_length= 128
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

def load_and_cache_examples( task, tokenizer, evaluate=False):

    processor = processors[task]()
    output_mode = output_modes[task]
    # Load data features from cache or dataset file
    cached_features_file = os.path.join(
        data_dir,
        "cached_{}_{}_{}_{}".format(
            "dev" if evaluate else "train",
            list(filter(None, model_name_or_path.split("/"))).pop(),
            str(max_seq_length),
            str(task),
        ),
    )
    if os.path.exists(cached_features_file):
        features = torch.load(cached_features_file)
    else:
        label_list = processor.get_labels()
        if task in ["mnli",
                    "mnli-mm"] and args.model_type in ["roberta",
                                                       "xlmroberta"]:
            # HACK(label indices are swapped in RoBERTa pretrained model)
            label_list[1], label_list[2] = label_list[2], label_list[1]
        examples = (
            processor.get_dev_examples(
                data_dir) if evaluate else processor.get_train_examples(
                data_dir))
        features = convert_examples_to_features(
            examples,
            tokenizer,
            max_length=max_seq_length,
            label_list=label_list,
            output_mode=output_mode,
        )
        torch.save(features, cached_features_file)

    # Convert to Tensors and build dataset
    all_input_ids = torch.tensor(
        [f.input_ids for f in features], dtype=torch.long)
    all_attention_mask = torch.tensor(
        [f.attention_mask for f in features], dtype=torch.long)
    all_token_type_ids = torch.tensor(
        [f.token_type_ids for f in features], dtype=torch.long)
    if output_mode == "classification":
        all_labels = torch.tensor(
            [f.label for f in features], dtype=torch.long)
    elif output_mode == "regression":
        all_labels = torch.tensor(
            [f.label for f in features], dtype=torch.float)

    dataset = TensorDataset(
        all_input_ids,
        all_attention_mask,
        all_token_type_ids,
        all_labels)
    return dataset

def evaluate(model, tokenizer, prefix=""):
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_task_names = ['qqp']
    eval_outputs_dirs = './tmp'
    results = {}
    for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):
        eval_dataset = load_and_cache_examples(
            eval_task, tokenizer, evaluate=True)

        if not os.path.exists(eval_output_dir):
            os.makedirs(eval_output_dir)

        # Note that DistributedSampler samples randomly
        eval_sampler = SequentialSampler(eval_dataset)
        eval_dataloader = DataLoader(
            eval_dataset,
            sampler=eval_sampler,
            batch_size=32)

        # Eval!

        eval_loss = 0.0
        nb_eval_steps = 0
        preds = None
        out_label_ids = None

        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.eval()
            batch = tuple(t.to(device) for t in batch)

            with torch.no_grad():
                inputs = {
                    "input_ids": batch[0],
                    "attention_mask": batch[1],
                    "labels": batch[3]}

                inputs["token_type_ids"] = batch[2]
                     # XLM, DistilBERT, RoBERTa, and XLM-RoBERTa don't use segment_ids
                # import pdb; pdb.set_trace()
                outputs = model(**inputs)
                tmp_eval_loss, logits = outputs[:2]

                eval_loss += tmp_eval_loss.mean().item()
            nb_eval_steps += 1
            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = inputs["labels"].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(
                    out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)

        eval_loss = eval_loss / nb_eval_steps

        from scipy.special import softmax

        probs = softmax(preds, axis=-1)
        entropy = np.exp((-probs * np.log(probs)).sum(axis=-1).mean())
        preds = np.argmax(preds, axis=1)

        result = compute_metrics(eval_task, preds, out_label_ids)
        results.update(result)
        if entropy is not None:
            result["eval_avg_entropy"] = entropy

        output_eval_file = os.path.join(
            eval_output_dir, prefix, "eval_results.txt")

    return results

from ModelVisual2 import ModelVisual
if __name__ == '__main__':
    prefix = 'bert_finegrained_0.95_onnx'
    os.makedirs('bert_finegrained_0.95_onnx', exist_ok=True)
    output_mode = output_modes[task_name]
    tokenizer = BertTokenizer.from_pretrained('textattack/bert-base-uncased-QQP')
    config = BertConfig.from_pretrained('textattack/bert-base-uncased-QQP')
    model = BertForSequenceClassification.from_pretrained('textattack/bert-base-uncased-QQP', config=config)
    dummy_input = torch.load('dummy_input.pth')
    data = (dummy_input['input_ids'], dummy_input['attention_mask'], dummy_input['token_type_ids'])
    state_dict = torch.load('./bert_finegrained_0.9_ck/bert_finegrained_0.95_nomask.pth')
    model.load_state_dict(state_dict)
    fine_grained_mask = {}
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            w_mask = torch.abs(module.weight.data)>0
            b_mask = torch.abs(module.bias.data)>0
            fine_grained_mask[name] = {'weight':w_mask, 'bias':b_mask}
            print(name, ' Sparsity ratio:', torch.sum(w_mask)/w_mask.numel())

    import pdb; pdb.set_trace()
    torch.onnx.export(model, data, os.path.join(prefix, 'bert_finegrained.onnx'), opset_version=10)

max_stars_repo_name: zheng-ningxin/MLPruning

max_issues_repo_name: zheng-ningxin/MLPruning

max_forks_repo_name: zheng-ningxin/MLPruning

reason: import numpy,from scipy


text:
import numpy as np
import scipy.sparse as sp
import torch
import scipy
import networkx as nx
from scipy import io

def encode_onehot(labels):
    classes = set(labels)
    classes_dict = {c: np.identity(len(classes))[i, :] for i, c in
                    enumerate(classes)}
    labels_onehot = np.array(list(map(classes_dict.get, labels)),
                             dtype=np.int32)
    return labels_onehot

def read_yelp_labels(file_tags):
    """
    Read labels of yelp dataset
    """
    X = np.loadtxt(file_tags)
    Y = np.int_(X)
    number_of_labels = Y.shape[1]
    for k in range(number_of_labels):
        if np.all((Y[:, k] == 0), axis=0):
            Y = np.delete(Y, k, 1)
    return Y

def read_flickr_labels(file_tags):
    features_struct = scipy.io.loadmat(file_tags)
    labels = scipy.sparse.csr_matrix(features_struct["group"])
    a = scipy.sparse.csr_matrix.toarray(labels)
    return a

def read_labels_nodes(file_tags):
    X = np.genfromtxt(file_tags, dtype=np.dtype(str))
    d_X = {"labels": X[:, -1], "nodes": X[:, 0]}
    return d_X

def new_load_data(G, sub_G, file_tags, num_of_nodes, nc_nodes, multilabel=False, features_file=None):
    nodes = list(sub_G.nodes())
    n2i = {nodes[i]: i for i in range(len(nodes))}
    common = list(set(nodes).intersection(set(nc_nodes)))
    train_nodes = list(set(nodes)-set(common))

    all_labels = {}
    with open(file_tags, 'r') as reader:
        for line in reader:
            parts = line.strip().split()
            n_id, l_id = str(parts[0]), str(parts[1])
            all_labels[n_id] = l_id
    reader.close()
    labels_d = {}
    for node in nodes:
        labels_d.update({node: all_labels[node]})
    labels = list(labels_d.values())
    labels = np.asarray(labels, dtype=np.dtype(str))
    labels = encode_onehot(labels)

    if features_file is None:
        features = sp.identity(num_of_nodes, dtype=np.float32)
    else:
        idx_features_labels = np.genfromtxt("../datasets/cora.content", dtype=np.dtype(str))
        features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)

    idx_map = {j: i for i, j in enumerate(nodes)}
    adj = sp.coo_matrix(nx.to_scipy_sparse_matrix(sub_G))

    # build symmetric adjacency matrix
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)

    features = normalize(features)
    adj = normalize(adj + sp.eye(adj.shape[0]))

    idx_train = [n2i[t] for t in train_nodes]
    idx_val = [n2i[t] for t in common]
    idx_test = [n2i[t] for t in common]

    features = torch.FloatTensor(np.array(features.todense()))
    if multilabel:
        labels = torch.FloatTensor(labels)
    else:
        labels = torch.LongTensor(np.where(labels)[1])

    adj = sparse_mx_to_torch_sparse_tensor(adj)

    idx_train = torch.LongTensor(idx_train)
    idx_val = torch.LongTensor(idx_val)
    idx_test = torch.LongTensor(idx_test)

    return adj, features, labels, idx_train, idx_val, idx_test

def normalize(mx):
    """Row-normalize sparse matrix"""
    rowsum = np.array(mx.sum(1))
    r_inv = np.power(rowsum, -1).flatten()
    r_inv[np.isinf(r_inv)] = 0.
    r_mat_inv = sp.diags(r_inv)
    mx = r_mat_inv.dot(mx)
    return mx

def accuracy(output, labels):
    preds = output.max(1)[1].type_as(labels)
    correct = preds.eq(labels).double()
    correct = correct.sum()
    return correct / len(labels)

def sparse_mx_to_torch_sparse_tensor(sparse_mx):
    """Convert a scipy sparse matrix to a torch sparse tensor."""
    sparse_mx = sparse_mx.tocoo().astype(np.float32)
    indices = torch.from_numpy(
        np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
    values = torch.from_numpy(sparse_mx.data)
    shape = torch.Size(sparse_mx.shape)
    return torch.sparse.FloatTensor(indices, values, shape)

ext: py

reason: import numpy,import scipy,from scipy,import networkx


text:
import numpy
from matplotlib import pyplot
from matplotlib import cm

## Variable declaration
# parameters
c = 1

# time
t_max = 0.5
nt = 151
dt = t_max/nt

# mesh
nx = 51
xmin = 0
xmax = 2
dx = (xmax-xmin)/(nx-1)
x = numpy.linspace(xmin, xmax, nx)

# Initialize dependent variable
u = numpy.ones(nx)
up = numpy.ones(nx)    # from previous time-step

# Initial condition: hat function
u[int(.5/dx):int(1/dx)+1]=2

def advance(t_steps):
    for n in range(0,t_steps): # loop over all time-steps
        up = u.copy()
        # with periodic BC all are internal points
        u[1:] = up[1:] - up[1:]*dt/dx*(up[1:] - up[0:-1])
        u[0] = up[0] - up[0]*dt/dx*(up[0] - up[-1])

def plot2d(x_plot, u_plot):
    fig = pyplot.figure(figsize=(5, 3), dpi=100)
    pyplot.xlim(0,2)
    pyplot.ylim(0,2.5)
    pyplot.plot(x_plot, u_plot)
    fig.show()

# check initial condition
#plot2d(x, u)

def plot_multi(x, u):
    fig = pyplot.figure(figsize=(7, 5), dpi=100)
    colour = iter(cm.rainbow(numpy.linspace(0,25,nt)))
    for i in range(0, nt, 25):
        c = next(colour)
        pyplot.plot(x,u, c=c)
        advance(25)
        pyplot.xlabel('x (m)')
        pyplot.ylabel('u (m/s)')
        pyplot.xlim(0, 2)
        pyplot.ylim(0, 2.5)
        pyplot.title('Non-linear convection')
    pyplot.show()

# Let us plot some curves to check if everything is alright
plot_multi(x, u)

reason: import numpy


text:
# This file will expose the player as a HTTP API

from flask import Flask, escape, request
from player import Player
import numpy as np

app = Flask(__name__)
player = Player(False)

@app.route('/health-check')
def health_check():
    return 'OK'

@app.route('/start', methods=['POST'])
def start():
    data = request.json
    state = np.asarray(data['state'])
    valid_actions = np.asarray(data['valid_actions'])
    action = player.start(state, valid_actions)
    return {
        'action': int(action)
    }

@app.route('/step', methods=['POST'])
def step():
    data = request.json
    state = np.asarray(data['state'])
    valid_actions = np.asarray(data['valid_actions'])
    prev_reward = data['prev_reward']
    action = player.step(state, valid_actions, prev_reward)
    return {
        'action': int(action)
    }

@app.route('/end', methods=['POST'])
def end():
    data = request.json
    state = np.asarray(data['state'])
    prev_reward = data['prev_reward']
    player.end(state, prev_reward)
    return {}

max_stars_repo_name: school-of-ai-angers/rl-arena

max_issues_repo_name: school-of-ai-angers/rl-arena

max_forks_repo_name: school-of-ai-angers/rl-arena


text:
import numpy as np
import joblib

from lightgbm import LGBMClassifier
from catboost import CatBoostClassifier
from xgboost import XGBClassifier

from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.ensemble import VotingClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier

from sklearn.metrics import accuracy_score
from sklearn.metrics import f1_score
from sklearn.metrics import confusion_matrix
from sklearn.pipeline import Pipeline
from sklearn.pipeline import make_pipeline
from sklearn.feature_selection import RFE

from sklearn.model_selection import cross_val_score
from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import cross_validate

SEED = 41

def load_model(**kwargs):
    if kwargs["model"] == "rf":
        return RandomForestClassifier(
            random_state=kwargs["random_state"],
            n_jobs=4,
            criterion="entropy",
            n_estimators=500
        )
    elif kwargs["model"] == "dt":
        return DecisionTreeClassifier(
            random_state=kwargs["random_state"],
            criterion="entropy",
            max_depth=8
        )
    elif kwargs["model"] == "lgb":
        return LGBMClassifier(
            random_state=kwargs["random_state"],
            objective="binary",
            metric="binary_logloss",
            class_weight="balanced",
            max_depth=30,
            num_leaves=1200
        )
    elif kwargs["model"] == "xgb":
        return XGBClassifier(
            random_state=kwargs["random_state"],
            use_label_encoder=False,
            objective="reg:logistic"
        )
    elif kwargs["model"] == "svm":
        return SVC(random_state=kwargs["random_state"])
    elif kwargs["model"] == "lr":
        return LogisticRegression(random_state=kwargs["random_state"], n_jobs=-1)
    elif kwargs["model"] == "knn":
        return KNeighborsClassifier(n_jobs=-1)
    elif kwargs["model"] == "adacat":
        bm = CatBoostClassifier(
            random_seed=kwargs["random_state"],
            logging_level="Silent",
            loss_function="Logloss",
            eval_metric="AUC",
            class_weights=(2.1, 1),
            thread_count=-1,
            early_stopping_rounds=10
        )
        return AdaBoostClassifier(
            base_estimator=bm,
            n_estimators=100,
            random_state=kwargs["random_state"]
        )
    elif kwargs["model"] == "adalgbm":
        bm = LGBMClassifier(
            random_state=kwargs["random_state"],
            objective="binary",
            metric="binary_logloss",
            class_weight="balanced",
            max_depth=30,
            num_leaves=1200
        )
        return AdaBoostClassifier(
            base_estimator=bm,
            n_estimators=100,
            random_state=kwargs["random_state"]
        )
    elif kwargs["model"] == "mlp":
        return MLPClassifier(
            random_state=kwargs["random_state"],
            activation="relu",
            solver="adam",
            batch_size=128,
            learning_rate_init=0.01,
            hidden_layer_sizes=(20, 12)
        )
    elif kwargs["model"] == "cat":
        return CatBoostClassifier(
            random_seed=kwargs["random_state"],
            logging_level="Silent",
            loss_function="Logloss",
            eval_metric="AUC",
            class_weights=(2.1, 1),
            thread_count=-1,
            early_stopping_rounds=10
        )
    elif kwargs["model"] == "voting":
        xgb = load_model(model="xgb", random_state=kwargs["random_state"])
        lgbm = load_model(model="lgb", random_state=kwargs["random_state"])
        cat = load_model(model="cat", random_state=kwargs["random_state"])
        adalgbm = load_model(model="adalgbm", random_state=kwargs["random_state"])
        adacat = load_model(model="adacat", random_state=kwargs["random_state"])
        voters = [
            ("xgb", xgb),
            ("lgbm", lgbm),
            ("catboost", cat),
            ("lgbmXadaboost", adalgbm),
            ("catboostXadaboost", adacat),
        ]
        return VotingClassifier(estimators=voters, voting="soft")
    else:
        print("Unsupported Algorithm")
        return None

def train(X_train, y_train, model):
    '''
        머신러닝 모델을 선택하여 학습을 진행하는 함수

        :param X_train: 학습할 2차원 리스트 특징벡터
        :param y_train: 학습할 1차원 리스트 레이블 벡터
        :param model: 문자열, 선택할 머신러닝 알고리즘
        :return: 학습된 머신러닝 모델 객체
    '''
    clf = load_model(model=model, random_state=SEED)
    clf.fit(X_train, y_train)

    # pipe = make_pipeline(
    #     MinMaxScaler(),
    #     clf
    # )
    # pipe.fit(X_train, y_train)
    return clf

def evaluate(X_test, y_test, model):
    '''
        학습된 머신러닝 모델로 검증 데이터를 검증하는 함수

        :param X_test: 검증할 2차원 리스트 특징 벡터
        :param y_test: 검증할 1차원 리스트 레이블 벡터
        :param model: 학습된 머신러닝 모델 객체
    '''
    predict = model.predict(X_test)
    print("정확도", model.score(X_test, y_test))

def evaluate_kfold(X_train, y_train, model):
    kfold = StratifiedKFold(n_splits=10).split(X_train, y_train)
    scores = []
    for k, (_train, _test) in enumerate(kfold):
        model.fit(X_train[_train], y_train[_train])
        local_prediction = model.predict(X_train[_test])
        score_acc = accuracy_score(local_prediction, y_train[_test])
        score_f1 = f1_score(local_prediction, y_train[_test])
        scores.append(score_acc)
        cf = confusion_matrix(local_prediction, y_train[_test])
        print(f"폴드: {k + 1}, 클래스 분포: {np.bincount(y_train[_test])}, 정확도: {score_acc:.4f}, f1: {score_f1:.4f}")
        print(cf)
    print("CV 정확도: %.4f +/- %.4f" % (np.mean(scores), np.std(scores)))

    # scores = cross_validate(
    #     estimator=model,
    #     X=X_train,
    #     y=y_train,
    #     scoring=["accuracy", "f1"],
    #     cv=10,
    #     n_jobs=-1,
    #     return_train_score=False
    # )

    # scores = cross_val_score(
    #     estimator=model,
    #     X=X_train,
    #     y=y_train,
    #     cv=10,
    #     n_jobs=1
    # )
    # print("CV 정확도 점수: %s" % scores)

def create_pretrained_model(X_train, y_train, model, path: str):
    trained = train(X_train, y_train, model)
    joblib.dump(trained, path)

def load_pretrained_model(path):
    return joblib.load(path)

max_stars_repo_name: team-missing-one/malware-predictor

max_issues_repo_name: team-missing-one/malware-predictor

max_forks_repo_name: team-missing-one/malware-predictor


text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import collections
import logging
import json
import math
import os
import random
import pickle
from tqdm import tqdm, trange
from os.path import join
from collections import Counter

import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler

from pytorch_pretrained_bert.tokenization import whitespace_tokenize, BasicTokenizer, BertTokenizer
from pytorch_pretrained_bert.modeling import BertModel
from pytorch_pretrained_bert.optimization import BertAdam
from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE

tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

class Example(object):

    def __init__(self,
                 qas_id,
                 qas_type,
                 doc_tokens,
                 question_text,
                 sent_num,
                 sent_names,
                 sup_fact_id,
                 sent_start_end_position,
                 entity_start_end_position,
                 orig_answer_text=None,
                 start_position=None,
                 end_position=None):
        self.qas_id = qas_id
        self.qas_type = qas_type
        self.doc_tokens = doc_tokens
        self.question_text = question_text
        self.sent_num = sent_num
        self.sent_names = sent_names
        self.sup_fact_id = sup_fact_id
        self.sent_start_end_position = sent_start_end_position
        self.entity_start_end_position = entity_start_end_position
        self.orig_answer_text = orig_answer_text
        self.start_position = start_position
        self.end_position = end_position

class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self,
                 qas_id,
                 doc_tokens,
                 doc_input_ids,
                 doc_input_mask,
                 doc_segment_ids,
                 query_tokens,
                 query_input_ids,
                 query_input_mask,
                 query_segment_ids,
                 sent_spans,
                 entity_spans,
                 sup_fact_ids,
                 token_to_orig_map,
                 start_position=None,
                 end_position=None):

        self.qas_id = qas_id
        self.doc_tokens = doc_tokens
        self.doc_input_ids = doc_input_ids
        self.doc_input_mask = doc_input_mask
        self.doc_segment_ids = doc_segment_ids

        self.query_tokens = query_tokens
        self.query_input_ids = query_input_ids
        self.query_input_mask = query_input_mask
        self.query_segment_ids = query_segment_ids

        self.sent_spans = sent_spans
        self.entity_spans = entity_spans
        self.sup_fact_ids = sup_fact_ids

        self.token_to_orig_map = token_to_orig_map

        self.start_position = start_position
        self.end_position = end_position

def clean_entity(entity):
    Type = entity[1]
    Text = entity[0]
    if Type == "DATE" and ',' in Text:
        Text = Text.replace(' ,', ',')
    if '?' in Text:
        Text = Text.split('?')[0]
    Text = Text.replace("\'\'", "\"")
    Text = Text.replace("# ", "#")
    return Text

def check_in_full_paras(answer, paras):
    full_doc = ""
    for p in paras:
        full_doc += " ".join(p[1])
    return answer in full_doc

def read_hotpot_examples(para_file, full_file, entity_file):
    with open(para_file, 'r', encoding='utf-8') as reader:
        para_data = json.load(reader)

    with open(full_file, 'r', encoding='utf-8') as reader:
        full_data = json.load(reader)

    with open(entity_file, 'r', encoding='utf-8') as reader:
        entity_data = json.load(reader)

    def is_whitespace(c):
        if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
            return True
        return False

    examples = []

    # for debug
    actually_in_case = 0
    failed_case = 0
    failed_sup_case = 0
    failed_case_paranum = []
    para_length = []
    bert_cutted_para_length = []
    for case in tqdm(full_data):
        key = case['_id']
        qas_type = case['type']
        sup_facts = set([(sp[0], sp[1])for sp in case['supporting_facts']])
        orig_answer_text = case['answer']

        sent_id = 0
        doc_tokens = []
        sent_names = []
        sup_facts_sent_id = []
        sent_start_end_position = []
        entity_start_end_position = []

        JUDGE_FLAG = orig_answer_text == 'yes' or orig_answer_text == 'no'
        FIND_FLAG = False

        char_to_word_offset = []  # Accumulated along all sentences
        prev_is_whitespace = True

        ans_start_position = None
        ans_end_position = None

        # for debug
        titles = set()

        for paragraph in para_data[key]:
            title = paragraph[0]
            sents = paragraph[1]
            if title in entity_data[key]:
                entities = entity_data[key][title]
            else:
                entities = []

            titles.add(title)

            for local_sent_id, sent in enumerate(sents):
                # Determine the global sent id for supporting facts
                local_sent_name = (title, local_sent_id)
                sent_names.append(local_sent_name)
                if local_sent_name in sup_facts:
                    sup_facts_sent_id.append(sent_id)
                sent_id += 1
                sent += " "

                sent_start_word_id = len(doc_tokens)
                sent_start_char_id = len(char_to_word_offset)

                for c in sent:
                    if is_whitespace(c):
                        prev_is_whitespace = True
                    else:
                        if prev_is_whitespace:
                            doc_tokens.append(c)
                        else:
                            doc_tokens[-1] += c
                        prev_is_whitespace = False
                    char_to_word_offset.append(len(doc_tokens) - 1)

                sent_end_word_id = len(doc_tokens) - 1
                sent_start_end_position.append((sent_start_word_id, sent_end_word_id))

                # Answer char position
                answer_offset = sent.find(orig_answer_text)
                if not JUDGE_FLAG and not FIND_FLAG and answer_offset != -1:
                    FIND_FLAG = True
                    start_char_position = sent_start_char_id + answer_offset
                    end_char_position = start_char_position + len(orig_answer_text) - 1
                    ans_start_position = char_to_word_offset[start_char_position]
                    ans_end_position = char_to_word_offset[end_char_position]

                # Find Entity Position
                entity_pointer = 0
                for entity in entities:
                    entity_text = clean_entity(entity)
                    entity_offset = sent.find(entity_text)
                    if entity_offset != -1:
                        entity_pointer += 1
                        start_char_position = sent_start_char_id + entity_offset
                        end_char_position = start_char_position + len(entity_text) - 1
                        ent_start_position = char_to_word_offset[start_char_position]
                        ent_end_position = char_to_word_offset[end_char_position]
                        entity_start_end_position.append((ent_start_position, ent_end_position, entity_text))
                    else:
                        break
                entities = entities[entity_pointer:]

                # Truncate longer document
                if len(doc_tokens) > 382:
                    break

            para_length.append(len(doc_tokens))

        # for ent in entity_start_end_position:
        #     print(ent)
        #     print(" ".join(doc_tokens[ent[0]:ent[1]+1]))
        # input()

        example = Example(
            qas_id=key,
            qas_type=qas_type,
            doc_tokens=doc_tokens,
            question_text=case['question'],
            sent_num=sent_id + 1,
            sent_names=sent_names,
            sup_fact_id=sup_facts_sent_id,
            sent_start_end_position=sent_start_end_position,
            entity_start_end_position=entity_start_end_position,
            orig_answer_text=orig_answer_text,
            start_position=ans_start_position,
            end_position=ans_end_position)
        examples.append(example)

    return examples

    #     # Check Finding Answer position
    #     if not JUDGE_FLAG:
    #         if end_position is None:
    #             failed_case += 1
    #             failed_case_paranum.append(len(para_data[key]))
    #             print(orig_answer_text)
    #             print(key)
    #             print("({}, {})".format(start_position, end_position))
    #             actual_text = " ".join(doc_tokens[start_position:(end_position + 1)])
    #             cleaned_answer_text = " ".join(whitespace_tokenize(orig_answer_text))
    #             if actual_text.find(cleaned_answer_text) == -1:
    #                 print("Could not find answer:")
    #             print("ACTUAL: ", actual_text)
    #             print("CLEANED: ", cleaned_answer_text)
    #             print(doc_tokens)
    #             input()
    #
    #     # Check finding supporting facts
    #     if len(sup_facts) != len(sup_facts_sent_id):
    #         failed_sup_case += 1
    #         print("gets:", titles)
    #         print("facts:", set([sp[0] for sp in sup_facts]))
    #         print("facts id:", sup_facts_sent_id)
    #
    # print("no answer: ", failed_case)
    # print("lack sup fact: ", failed_sup_case)
    # print("total cases: ", len(full_data))
    # print("noanswer para num: ", Counter(failed_case_paranum))
    # print("Avg bert len: ", np.average(np.array(bert_cutted_para_length)))
    # print("Max bert len: ", np.max(np.array(bert_cutted_para_length)))
    # np.array(bert_cutted_para_length).dump("bert_doc_len_aug.np")

def convert_examples_to_features(examples, tokenizer, max_seq_length, max_query_length):
    features = []
    failed = 0
    ans_failed = 0
    for (example_index, example) in enumerate(tqdm(examples)):
        query_tokens = ["[CLS]"] + tokenizer.tokenize(example.question_text)
        if len(query_tokens) > max_query_length - 1:
            query_tokens = query_tokens[:max_query_length - 1]
        query_tokens.append("[SEP]")

        tok_to_orig_index = [0]
        orig_to_tok_index = []
        orig_to_tok_back_index = []
        all_doc_tokens = ["[CLS]"]
        sentence_spans = []
        entity_spans = []
        for (i, token) in enumerate(example.doc_tokens):
            orig_to_tok_index.append(len(all_doc_tokens))
            sub_tokens = tokenizer.tokenize(token)
            for sub_token in sub_tokens:
                tok_to_orig_index.append(i)
                all_doc_tokens.append(sub_token)
            orig_to_tok_back_index.append(len(all_doc_tokens) - 1)

        def relocate_tok_span(orig_start_position, orig_end_position, orig_text):
            if orig_start_position is None:
                return 0, 0

            global tokenizer
            nonlocal orig_to_tok_index, example, all_doc_tokens

            tok_start_position = orig_to_tok_index[orig_start_position]
            if orig_end_position < len(example.doc_tokens) - 1:
                tok_end_position = orig_to_tok_index[orig_end_position + 1] - 1
            else:
                tok_end_position = len(all_doc_tokens) - 1
            # Make answer span more accurate.
            return _improve_answer_span(
                all_doc_tokens, tok_start_position, tok_end_position, tokenizer, orig_text)

        ans_start_position, ans_end_position \
            = relocate_tok_span(example.start_position, example.end_position, example.orig_answer_text)

        for entity_span in example.entity_start_end_position:
            ent_start_position, ent_end_position \
                = relocate_tok_span(entity_span[0], entity_span[1], entity_span[2])
            entity_spans.append((ent_start_position, ent_end_position, entity_span[2]))

        for sent_span in example.sent_start_end_position:
            if sent_span[0] >= len(orig_to_tok_index) or sent_span[0] >= sent_span[1]:
                continue
            sent_start_position = orig_to_tok_index[sent_span[0]]
            sent_end_position = orig_to_tok_back_index[sent_span[1]]
            sentence_spans.append((sent_start_position, sent_end_position))

        # Padding Document
        all_doc_tokens = all_doc_tokens[:max_seq_length - 1] + ["[SEP]"]
        doc_input_ids = tokenizer.convert_tokens_to_ids(all_doc_tokens)
        doc_input_mask = [1] * len(doc_input_ids)
        doc_segment_ids = [0] * len(doc_input_ids)

        while len(doc_input_ids) < max_seq_length:
            doc_input_ids.append(0)
            doc_input_mask.append(0)
            doc_segment_ids.append(0)

        # Padding Question
        query_input_ids = tokenizer.convert_tokens_to_ids(query_tokens)
        query_input_mask = [1] * len(query_input_ids)
        query_segment_ids = [0] * len(query_input_ids)

        while len(query_input_ids) < max_query_length:
            query_input_ids.append(0)
            query_input_mask.append(0)
            query_segment_ids.append(0)

        assert len(doc_input_ids) == max_seq_length
        assert len(doc_input_mask) == max_seq_length
        assert len(doc_segment_ids) == max_seq_length
        assert len(query_input_ids) == max_query_length
        assert len(query_input_mask) == max_query_length
        assert len(query_segment_ids) == max_query_length

        # Dropout out-of-bound span
        entity_spans = entity_spans[:_largest_valid_index(entity_spans, max_seq_length)]
        sentence_spans = sentence_spans[:_largest_valid_index(sentence_spans, max_seq_length)]

        sup_fact_ids = example.sup_fact_id
        sent_num = len(sentence_spans)
        sup_fact_ids = [sent_id for sent_id in sup_fact_ids if sent_id < sent_num]
        if len(sup_fact_ids) != len(example.sup_fact_id):
            failed += 1
        if ans_start_position >= 512:
            ans_failed += 1

        # print("ALL:\n", all_doc_tokens)
        # print("MASK:\n", doc_input_mask)
        # print("ANS:\n", all_doc_tokens[ans_start_position: ans_end_position + 1])
        # print("SP_FACTS: \n", sup_fact_ids)
        # print("ANSWER:\n", example.orig_answer_text)
        # print("ANSWER:\n", all_doc_tokens[ans_start_position: ans_end_position + 1])
        # for i, sent in enumerate(sentence_spans):
        #     print("sent{}:\n".format(i), all_doc_tokens[sent[0]: sent[1] + 1])
        #     os, oe = tok_to_orig_index[sent[0]], tok_to_orig_index[sent[1]]
        #     print("ORI_SENT:\n", example.doc_tokens[os: oe + 1])
        #     input()
        #     for ent in entity_spans:
        #         if ent[0] >= sent[0] and ent[1] <= sent[1]:
        #             print("ORI:  ", ent[2])
        #             print("NEW:  ", all_doc_tokens[ent[0] : ent[1] + 1])
        #     input()

        features.append(
            InputFeatures(qas_id=example.qas_id,
                          doc_tokens=all_doc_tokens,
                          doc_input_ids=doc_input_ids,
                          doc_input_mask=doc_input_mask,
                          doc_segment_ids=doc_segment_ids,
                          query_tokens=query_tokens,
                          query_input_ids=query_input_ids,
                          query_input_mask=query_input_mask,
                          query_segment_ids=query_segment_ids,
                          sent_spans=sentence_spans,
                          entity_spans=entity_spans,
                          sup_fact_ids=sup_fact_ids,
                          token_to_orig_map=tok_to_orig_index,
                          start_position=ans_start_position,
                          end_position=ans_end_position)
        )
    # print("Failed: ", failed)
    # print("Ans_Failed: ", ans_failed)
    return features

def _largest_valid_index(spans, limit):
    for idx in range(len(spans)):
        if spans[idx][1] >= limit:
            return idx

def _improve_answer_span(doc_tokens, input_start, input_end, tokenizer,
                         orig_answer_text):
    """Returns tokenized answer spans that better match the annotated answer."""

    tok_answer_text = " ".join(tokenizer.tokenize(orig_answer_text))

    for new_start in range(input_start, input_end + 1):
        for new_end in range(input_end, new_start - 1, -1):
            text_span = " ".join(doc_tokens[new_start:(new_end + 1)])
            if text_span == tok_answer_text:
                return new_start, new_end

    return input_start, input_end

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    ## Required parameters
    parser.add_argument("--bert_model", default='bert-base-cased', type=str)

    ## Other parameters
    parser.add_argument("--do_lower_case", default=False, action='store_true',
                        help="Set this flag if you are using an uncased model.")
    parser.add_argument("--layers", default="-1,-2,-3", type=str)
    parser.add_argument("--max_seq_length", default=512, type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. Sequences longer "
                             "than this will be truncated, and sequences shorter than this will be padded.")
    parser.add_argument("--batch_size", default=15, type=int, help="Batch size for predictions.")

    args = parser.parse_args()

    data_path = '/media/disk2/jennybae/HotpotQA/'
    dev_para_path = join(data_path, 'Selected_Paras', 'dev_paras.json')
    dev_full_path = join(data_path, 'hotpot_dev_distractor_v1.json')
    dev_entity_path = join(data_path, 'Selected_Paras', 'dev_entity.json')
    # examples = read_hotpot_examples(para_file=dev_para_path, full_file=dev_full_path, entity_file=dev_entity_path)
    # examples = pickle.load(open(join(data_path, 'BERT_Features', 'examples', 'dev_example.pkl'), 'rb'))
    # pickle.dump(examples, open(join(data_path, 'BERT_Features', 'examples', 'dev_example.pkl'), 'wb'))
    # features = convert_examples_to_features(examples, tokenizer, max_seq_length=512, max_query_length=50)
    # pickle.dump(features, open(join(data_path, 'BERT_Features', 'features', 'dev_feature.pkl'), 'wb'))
    features = pickle.load(open(join(data_path, 'BERT_Features', 'features', 'dev_feature.pkl'), 'rb'))

    model = BertModel.from_pretrained(args.bert_model)
    model.cuda()
    model = torch.nn.DataParallel(model)
    model.eval()

    doc_input_ids = torch.LongTensor([f.doc_input_ids for f in features])
    doc_input_mask = torch.LongTensor([f.doc_input_mask for f in features])
    query_input_ids = torch.LongTensor([f.query_input_ids for f in features])
    query_input_mask = torch.LongTensor([f.query_input_mask for f in features])
    all_example_indices = torch.arange(doc_input_ids.size(0), dtype=torch.long)

    eval_data = TensorDataset(doc_input_ids, doc_input_mask, query_input_ids, query_input_mask, all_example_indices)
    eval_sampler = SequentialSampler(eval_data)
    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.batch_size)

    layer_indices = [int(x) for x in args.layers.split(",")]

    output_json = dict()
    out_feature_dir = join(data_path, 'BERT_Features', 'layers')
    file_id = 0
    for doc_ids, doc_mask, query_ids, query_mask, example_indices in tqdm(eval_dataloader):
        doc_ids = doc_ids.cuda()
        doc_mask = doc_mask.cuda()

        all_doc_encoder_layers, _ = model(doc_ids, token_type_ids=None, attention_mask=doc_mask)
        all_query_encoder_layers, _ = model(query_ids, token_type_ids=None, attention_mask=query_mask)
        selected_doc_layers = [all_doc_encoder_layers[layer_index].detach().cpu().numpy() for layer_index in layer_indices]
        selected_query_layers = [all_query_encoder_layers[layer_index].detach().cpu().numpy() for layer_index in layer_indices]

        for b, example_index in enumerate(example_indices):
            feature = features[example_index.item()]
            case = dict()
            case['query'] = [layer[b] for layer in selected_query_layers]
            case['doc'] = [layer[b] for layer in selected_doc_layers]
            output_json[feature.qas_id] = case

    output_file = join(out_feature_dir, "dev_layers.pkl")
    pickle.dump(output_json, open(output_file, 'wb'))
    output_json = dict()
    file_id += 1