text:
# -*- coding: utf-8 -*-
# @Author: Jie
# @Date:   2017-06-15 14:11:08
# @Last Modified by:   Jie Yang,     Contact: jieynlp@gmail.com
# @Last Modified time: 2019-01-14 16:09:16

from __future__ import print_function
import time
import sys
import argparse
import random
import torch
import gc
import logging
import torch.optim as optim
import numpy as np
from utils.metric import get_ner_fmeasure
from model.seqlabel import SeqLabel
from model.sentclassifier import SentClassifier
from utils.data import Data
import utils.utils as utils

try:
    import cPickle as pickle
except ImportError:
    import pickle

seed_num = 42
random.seed(seed_num)
torch.manual_seed(seed_num)
np.random.seed(seed_num)

def data_initialization(data):
    data.initial_feature_alphabets()
    data.build_alphabet(data.train_dir)
    data.build_alphabet(data.dev_dir)
    data.build_alphabet(data.test_dir)
    data.fix_alphabet()

def predict_check(pred_variable,
                  gold_variable,
                  mask_variable,
                  sentence_classification=False):
    """
        input:
            pred_variable (batch_size, sent_len): pred tag result, in numpy format
            gold_variable (batch_size, sent_len): gold result variable
            mask_variable (batch_size, sent_len): mask variable
    """
    pred = pred_variable.cpu().data.numpy()
    gold = gold_variable.cpu().data.numpy()
    mask = mask_variable.cpu().data.numpy()
    overlaped = (pred == gold)
    if sentence_classification:
        right_token = np.sum(overlaped)
        total_token = overlaped.shape[0]  # =batch_size
    else:
        right_token = np.sum(overlaped * mask)
        total_token = mask.sum()
    # logging.info("right: %s, total: %s"%(right_token, total_token))
    return right_token, total_token

def recover_label(pred_variable,
                  gold_variable,
                  mask_variable,
                  label_alphabet,
                  word_recover,
                  sentence_classification=False):
    """
        input:
            pred_variable (batch_size, sent_len): pred tag result
            gold_variable (batch_size, sent_len): gold result variable
            mask_variable (batch_size, sent_len): mask variable
    """
    pred_variable = pred_variable[word_recover]
    gold_variable = gold_variable[word_recover]
    mask_variable = mask_variable[word_recover]
    batch_size = gold_variable.size(0)
    if sentence_classification:
        pred_tag = pred_variable.cpu().data.numpy().tolist()
        gold_tag = gold_variable.cpu().data.numpy().tolist()
        pred_label = [label_alphabet.get_instance(pred) for pred in pred_tag]
        gold_label = [label_alphabet.get_instance(gold) for gold in gold_tag]
    else:
        seq_len = gold_variable.size(1)
        mask = mask_variable.cpu().data.numpy()
        pred_tag = pred_variable.cpu().data.numpy()
        gold_tag = gold_variable.cpu().data.numpy()
        batch_size = mask.shape[0]
        pred_label = []
        gold_label = []
        for idx in range(batch_size):
            pred = [
                label_alphabet.get_instance(pred_tag[idx][idy])
                for idy in range(seq_len) if mask[idx][idy] != 0
            ]
            gold = [
                label_alphabet.get_instance(gold_tag[idx][idy])
                for idy in range(seq_len) if mask[idx][idy] != 0
            ]
            assert (len(pred) == len(gold))
            pred_label.append(pred)
            gold_label.append(gold)
    return pred_label, gold_label

def recover_nbest_label(pred_variable, mask_variable, label_alphabet,
                        word_recover):
    """
        input:
            pred_variable (batch_size, sent_len, nbest): pred tag result
            mask_variable (batch_size, sent_len): mask variable
            word_recover (batch_size)
        output:
            nbest_pred_label list: [batch_size, nbest, each_seq_len]
    """
    # exit(0)
    pred_variable = pred_variable[word_recover]
    mask_variable = mask_variable[word_recover]
    batch_size = pred_variable.size(0)
    seq_len = pred_variable.size(1)
    nbest = pred_variable.size(2)
    mask = mask_variable.cpu().data.numpy()
    pred_tag = pred_variable.cpu().data.numpy()
    batch_size = mask.shape[0]
    pred_label = []
    for idx in range(batch_size):
        pred = []
        for idz in range(nbest):
            each_pred = [
                label_alphabet.get_instance(pred_tag[idx][idy][idz])
                for idy in range(seq_len) if mask[idx][idy] != 0
            ]
            pred.append(each_pred)
        pred_label.append(pred)
    return pred_label

def lr_decay(optimizer, epoch, decay_rate, init_lr):
    """
    学习率衰减
    :param optimizer:
    :param epoch:
    :param decay_rate:
    :param init_lr:
    :return:
    """
    lr = init_lr / (1 + decay_rate * epoch)
    logging.info(" Learning rate is set as: %s" % lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return optimizer

def evaluate(data, model, name, nbest=None):
    if name == "train":
        instances = data.train_Ids
    elif name == "dev":
        instances = data.dev_Ids
    elif name == 'test':
        instances = data.test_Ids
    elif name == 'raw':
        instances = data.raw_Ids
    else:
        logging.info("Error: wrong evaluate name, %s" % name)
        exit(1)
    right_token = 0
    whole_token = 0
    nbest_pred_results = []
    pred_scores = []
    pred_results = []
    gold_results = []
    # set model in eval model
    model.eval()
    batch_size = data.HP_batch_size
    start_time = time.time()
    train_num = len(instances)
    total_batch = train_num // batch_size + 1
    for batch_id in range(total_batch):
        start = batch_id * batch_size
        end = (batch_id + 1) * batch_size
        if end > train_num:
            end = train_num
        instance = instances[start:end]
        if not instance:
            continue
        batch_word, batch_features, batch_wordlen, batch_wordrecover, batch_char, \
        batch_charlen, batch_charrecover, batch_label, mask = batchify_with_label(instance,
                                                                                data.HP_gpu,
                                                                                False,
                                                                                data.sentence_classification)
        if nbest and not data.sentence_classification:
            scores, nbest_tag_seq = model.decode_nbest(
                batch_word, batch_features, batch_wordlen, batch_char,
                batch_charlen, batch_charrecover, mask, nbest)
            nbest_pred_result = recover_nbest_label(
                nbest_tag_seq, mask, data.label_alphabet, batch_wordrecover)
            nbest_pred_results += nbest_pred_result
            pred_scores += scores[batch_wordrecover].cpu().data.numpy().tolist(
            )
            # select the best sequence to evalurate
            tag_seq = nbest_tag_seq[:, :, 0]
        else:
            tag_seq = model(batch_word, batch_features, batch_wordlen,
                            batch_char, batch_charlen, batch_charrecover, mask)
        # logging.info("tag:",tag_seq)
        pred_label, gold_label = recover_label(
            tag_seq, batch_label, mask, data.label_alphabet, batch_wordrecover,
            data.sentence_classification)
        pred_results += pred_label
        gold_results += gold_label
    decode_time = time.time() - start_time
    speed = len(instances) / decode_time
    acc, p, r, f = get_ner_fmeasure(gold_results, pred_results, data.tagScheme)
    if nbest and not data.sentence_classification:
        return speed, acc, p, r, f, nbest_pred_results, pred_scores
    return speed, acc, p, r, f, pred_results, pred_scores

def batchify_with_label(input_batch_list,
                        gpu,
                        if_train=True,
                        sentence_classification=False):
    if sentence_classification:
        return batchify_sentence_classification_with_label(
            input_batch_list, gpu, if_train)
    else:
        return batchify_sequence_labeling_with_label(input_batch_list, gpu,
                                                     if_train)

def batchify_sequence_labeling_with_label(input_batch_list, gpu,
                                          if_train=True):
    """
    每个batch内部按照最长的语句做padding操作
    input: list of words, chars and labels, various length. [[words, features, chars, labels]...]
        从input_batch_list里面很容易转化为下面的格式
        words: word ids for one sentence. (batch_size, sent_len)
        features: features ids for one sentence. (batch_size, sent_len, feature_num)
        chars: char ids for on sentences, various length. (batch_size, sent_len, each_word_length)
        labels: label ids for one sentence. (batch_size, sent_len)

    output:
        zero padding for word and char, with their batch length
        word_seq_tensor: (batch_size, max_sent_len) Variable
        feature_seq_tensors: [(batch_size, max_sent_len),...] list of Variable
        word_seq_lengths: (batch_size,1) Tensor
        char_seq_tensor: (batch_size*max_sent_len, max_word_len) Variable
        char_seq_lengths: (batch_size*max_sent_len,1) Tensor
        char_seq_recover: (batch_size*max_sent_len,1) recover char sequence order
        label_seq_tensor: (batch_size, max_sent_len)
        mask: (batch_size, max_sent_len)
    """

    # 从input_batch_list里面转化为下面的格式，第一维为batch_size
    # logging.info(input_batch_list)
    batch_size = len(input_batch_list)
    words = [sent[0] for sent in input_batch_list]
    features = [np.asarray(sent[1]) for sent in input_batch_list]
    feature_num = len(features[0][0])
    chars = [sent[2] for sent in input_batch_list]
    labels = [sent[3] for sent in input_batch_list]

    word_seq_lengths = torch.LongTensor(list(map(len, words)))
    max_seq_len = word_seq_lengths.max().item()
    word_seq_tensor = torch.zeros((batch_size, max_seq_len),
                                  requires_grad=if_train).long()
    label_seq_tensor = torch.zeros((batch_size, max_seq_len),
                                   requires_grad=if_train).long()
    feature_seq_tensors = []
    for idx in range(feature_num):
        feature_seq_tensors.append(
            torch.zeros((batch_size, max_seq_len),
                        requires_grad=if_train).long())
    mask = torch.zeros((batch_size, max_seq_len),
                       requires_grad=if_train).byte()
    for idx, (seq, label, seqlen) in enumerate(
            zip(words, labels, word_seq_lengths)):
        seqlen = seqlen.item()
        word_seq_tensor[idx, :seqlen] = torch.LongTensor(seq)
        label_seq_tensor[idx, :seqlen] = torch.LongTensor(label)
        mask[idx, :seqlen] = torch.Tensor([1] * seqlen)
        for idy in range(feature_num):
            feature_seq_tensors[idy][idx, :seqlen] = torch.LongTensor(
                features[idx][:, idy])
    word_seq_lengths, word_perm_idx = word_seq_lengths.sort(0, descending=True)
    word_seq_tensor = word_seq_tensor[word_perm_idx]
    for idx in range(feature_num):
        feature_seq_tensors[idx] = feature_seq_tensors[idx][word_perm_idx]
    label_seq_tensor = label_seq_tensor[word_perm_idx]
    mask = mask[word_perm_idx]

    # deal with char
    # pad_chars (batch_size, max_seq_len)
    pad_chars = [
        chars[idx] + [[0]] * (max_seq_len - len(chars[idx]))
        for idx in range(len(chars))
    ]
    length_list = [list(map(len, pad_char)) for pad_char in pad_chars]
    max_word_len = max(map(max, length_list))
    char_seq_tensor = torch.zeros((batch_size, max_seq_len, max_word_len),
                                  requires_grad=if_train).long()
    char_seq_lengths = torch.LongTensor(length_list)
    for idx, (seq, seqlen) in enumerate(zip(pad_chars, char_seq_lengths)):
        for idy, (word, wordlen) in enumerate(zip(seq, seqlen)):
            # print len(word), wordlen
            char_seq_tensor[idx, idy, :wordlen] = torch.LongTensor(word)

    char_seq_tensor = char_seq_tensor[word_perm_idx].view(
        batch_size * max_seq_len, -1)
    char_seq_lengths = char_seq_lengths[word_perm_idx].view(
        batch_size * max_seq_len, )
    char_seq_lengths, char_perm_idx = char_seq_lengths.sort(0, descending=True)
    char_seq_tensor = char_seq_tensor[char_perm_idx]
    _, char_seq_recover = char_perm_idx.sort(0, descending=False)
    _, word_seq_recover = word_perm_idx.sort(0, descending=False)
    if gpu:
        word_seq_tensor = word_seq_tensor.cuda()
        for idx in range(feature_num):
            feature_seq_tensors[idx] = feature_seq_tensors[idx].cuda()
        word_seq_lengths = word_seq_lengths.cuda()
        word_seq_recover = word_seq_recover.cuda()
        label_seq_tensor = label_seq_tensor.cuda()
        char_seq_tensor = char_seq_tensor.cuda()
        char_seq_recover = char_seq_recover.cuda()
        mask = mask.cuda()
    # logging.info(word_seq_tensor)
    # logging.info(feature_seq_tensors)
    # logging.info(word_seq_lengths)
    # logging.info(word_seq_recover)
    # logging.info(char_seq_tensor)
    # logging.info(char_seq_lengths)
    # logging.info(char_seq_recover)
    # logging.info(label_seq_tensor)
    # logging.info(mask)
    return word_seq_tensor, feature_seq_tensors, word_seq_lengths, word_seq_recover, char_seq_tensor, \
           char_seq_lengths, char_seq_recover, label_seq_tensor, mask

def batchify_sentence_classification_with_label(input_batch_list,
                                                gpu,
                                                if_train=True):
    """
        input: list of words, chars and labels, various length. [[words, features, chars, labels]...]
            words: word ids for one sentence. (batch_size, sent_len)
            features: features ids for one sentence. (batch_size, feature_num), each sentence has one set of feature
            chars: char ids for on sentences, various length. (batch_size, sent_len, each_word_length)
            labels: label ids for one sentence. (batch_size,), each sentence has one set of feature

        output:
            zero padding for word and char, with their batch length
            word_seq_tensor: (batch_size, max_sent_len) Variable
            feature_seq_tensors: [(batch_size,), ... ] list of Variable
            word_seq_lengths: (batch_size,1) Tensor
            char_seq_tensor: (batch_size*max_sent_len, max_word_len) Variable
            char_seq_lengths: (batch_size*max_sent_len,1) Tensor
            char_seq_recover: (batch_size*max_sent_len,1)  recover char sequence order
            label_seq_tensor: (batch_size, )
            mask: (batch_size, max_sent_len)
    """

    batch_size = len(input_batch_list)
    words = [sent[0] for sent in input_batch_list]
    features = [np.asarray(sent[1]) for sent in input_batch_list]
    feature_num = len(features[0])
    chars = [sent[2] for sent in input_batch_list]
    labels = [sent[3] for sent in input_batch_list]
    word_seq_lengths = torch.LongTensor(list(map(len, words)))
    max_seq_len = word_seq_lengths.max().item()
    word_seq_tensor = torch.zeros((batch_size, max_seq_len),
                                  requires_grad=if_train).long()
    label_seq_tensor = torch.zeros((batch_size, ),
                                   requires_grad=if_train).long()
    feature_seq_tensors = []
    for idx in range(feature_num):
        feature_seq_tensors.append(
            torch.zeros((batch_size, max_seq_len),
                        requires_grad=if_train).long())
    mask = torch.zeros((batch_size, max_seq_len),
                       requires_grad=if_train).byte()
    label_seq_tensor = torch.LongTensor(labels)
    # exit(0)
    for idx, (seq, seqlen) in enumerate(zip(words, word_seq_lengths)):
        seqlen = seqlen.item()
        word_seq_tensor[idx, :seqlen] = torch.LongTensor(seq)
        mask[idx, :seqlen] = torch.Tensor([1] * seqlen)
        for idy in range(feature_num):
            feature_seq_tensors[idy][idx, :seqlen] = torch.LongTensor(
                features[idx][:, idy])
    word_seq_lengths, word_perm_idx = word_seq_lengths.sort(0, descending=True)
    word_seq_tensor = word_seq_tensor[word_perm_idx]
    for idx in range(feature_num):
        feature_seq_tensors[idx] = feature_seq_tensors[idx][word_perm_idx]
    label_seq_tensor = label_seq_tensor[word_perm_idx]
    mask = mask[word_perm_idx]
    ## deal with char
    # pad_chars (batch_size, max_seq_len)
    pad_chars = [
        chars[idx] + [[0]] * (max_seq_len - len(chars[idx]))
        for idx in range(len(chars))
    ]
    length_list = [list(map(len, pad_char)) for pad_char in pad_chars]
    max_word_len = max(map(max, length_list))
    char_seq_tensor = torch.zeros((batch_size, max_seq_len, max_word_len),
                                  requires_grad=if_train).long()
    char_seq_lengths = torch.LongTensor(length_list)
    for idx, (seq, seqlen) in enumerate(zip(pad_chars, char_seq_lengths)):
        for idy, (word, wordlen) in enumerate(zip(seq, seqlen)):
            # print len(word), wordlen
            char_seq_tensor[idx, idy, :wordlen] = torch.LongTensor(word)

    char_seq_tensor = char_seq_tensor[word_perm_idx].view(
        batch_size * max_seq_len, -1)
    char_seq_lengths = char_seq_lengths[word_perm_idx].view(
        batch_size * max_seq_len, )
    char_seq_lengths, char_perm_idx = char_seq_lengths.sort(0, descending=True)
    char_seq_tensor = char_seq_tensor[char_perm_idx]
    _, char_seq_recover = char_perm_idx.sort(0, descending=False)
    _, word_seq_recover = word_perm_idx.sort(0, descending=False)
    if gpu:
        word_seq_tensor = word_seq_tensor.cuda()
        for idx in range(feature_num):
            feature_seq_tensors[idx] = feature_seq_tensors[idx].cuda()
        word_seq_lengths = word_seq_lengths.cuda()
        word_seq_recover = word_seq_recover.cuda()
        label_seq_tensor = label_seq_tensor.cuda()
        char_seq_tensor = char_seq_tensor.cuda()
        char_seq_recover = char_seq_recover.cuda()
        mask = mask.cuda()
    return word_seq_tensor, feature_seq_tensors, word_seq_lengths, word_seq_recover, char_seq_tensor, \
           char_seq_lengths, char_seq_recover, label_seq_tensor, mask

def train(data):
    logging.info("Training model...")
    data.show_data_summary()
    save_data_name = data.model_dir + ".dset"
    data.save(save_data_name)
    if data.sentence_classification:
        model = SentClassifier(data)
    else:
        model = SeqLabel(data)

    # loss_function = nn.NLLLoss()
    if data.optimizer.lower() == "sgd":
        optimizer = optim.SGD(
            model.parameters(),
            lr=data.HP_lr,
            momentum=data.HP_momentum,
            weight_decay=data.HP_l2)
    elif data.optimizer.lower() == "adagrad":
        optimizer = optim.Adagrad(
            model.parameters(), lr=data.HP_lr, weight_decay=data.HP_l2)
    elif data.optimizer.lower() == "adadelta":
        optimizer = optim.Adadelta(
            model.parameters(), lr=data.HP_lr, weight_decay=data.HP_l2)
    elif data.optimizer.lower() == "rmsprop":
        optimizer = optim.RMSprop(
            model.parameters(), lr=data.HP_lr, weight_decay=data.HP_l2)
    elif data.optimizer.lower() == "adam":
        optimizer = optim.Adam(
            model.parameters(), lr=data.HP_lr, weight_decay=data.HP_l2)
    else:
        logging.info("Optimizer illegal: %s" % (data.optimizer))
        exit(1)

    best_dev = -10
    # start training
    for idx in range(data.HP_iteration):
        epoch_start = time.time()
        temp_start = epoch_start
        logging.info("Epoch: %s/%s" % (idx, data.HP_iteration))
        if data.optimizer == "SGD":
            optimizer = lr_decay(optimizer, idx, data.HP_lr_decay, data.HP_lr)
        instance_count = 0
        sample_id = 0
        sample_loss = 0
        total_loss = 0
        right_token = 0
        whole_token = 0
        random.shuffle(data.train_Ids)
        logging.info(
            "Shuffle: first input word list:%s" % data.train_Ids[0][0])

        # set model in train model
        model.train()
        model.zero_grad()
        batch_size = data.HP_batch_size
        batch_id = 0
        train_num = len(data.train_Ids)
        total_batch = train_num // batch_size + 1
        for batch_id in range(total_batch):
            start = batch_id * batch_size
            end = (batch_id + 1) * batch_size
            if end > train_num:
                end = train_num

            # 一个batch内的input
            instance = data.train_Ids[start:end]
            if not instance:
                continue
            batch_word, batch_features, batch_wordlen, batch_wordrecover, batch_char, batch_charlen, \
            batch_charrecover, batch_label, mask = batchify_with_label(instance,
                                                                       data.HP_gpu,
                                                                       True,
                                                                       data.sentence_classification)
            instance_count += 1
            loss, tag_seq = model.neg_log_likelihood_loss(
                batch_word, batch_features, batch_wordlen, batch_char,
                batch_charlen, batch_charrecover, batch_label, mask)
            right, whole = predict_check(tag_seq, batch_label, mask,
                                         data.sentence_classification)
            right_token += right
            whole_token += whole
            # logging.info("loss:",loss.item())
            sample_loss += loss.item()
            total_loss += loss.item()
            if end % 50000 == 0:
                temp_time = time.time()
                temp_cost = temp_time - temp_start
                temp_start = temp_time
                logging.info(
                    "     Instance: %s; Time: %.2fs; loss: %.4f; acc: %s/%s=%.4f"
                    % (end, temp_cost, sample_loss, right_token, whole_token,
                       (right_token + 0.) / whole_token))
                if sample_loss > 1e8 or str(sample_loss) == "nan":
                    logging.info(
                        "ERROR: LOSS EXPLOSION (>1e8) ! PLEASE SET PROPER PARAMETERS AND STRUCTURE! EXIT...."
                    )
                    exit(1)
                sys.stdout.flush()
                sample_loss = 0
            loss.backward()
            optimizer.step()
            model.zero_grad()
        temp_time = time.time()
        temp_cost = temp_time - temp_start
        logging.info(
            "     Instance: %s; Time: %.2fs; loss: %.4f; acc: %s/%s=%.4f" %
            (end, temp_cost, sample_loss, right_token, whole_token,
             (right_token + 0.) / whole_token))

        epoch_finish = time.time()
        epoch_cost = epoch_finish - epoch_start
        logging.info(
            "Epoch: %s training finished. Time: %.2fs, speed: %.2fst/s,  total loss: %s"
            % (idx, epoch_cost, train_num / epoch_cost, total_loss))
        logging.info("totalloss: %s" % total_loss)
        if total_loss > 1e8 or str(total_loss) == "nan":
            logging.info(
                "ERROR: LOSS EXPLOSION (>1e8) ! PLEASE SET PROPER PARAMETERS AND STRUCTURE! EXIT...."
            )
            exit(1)
        # continue
        speed, acc, p, r, f, _, _ = evaluate(data, model, "dev")
        dev_finish = time.time()
        dev_cost = dev_finish - epoch_finish

        if data.seg:
            current_score = f
            logging.info(
                "Dev: time: %.2fs, speed: %.2fst/s; acc: %.4f, p: %.4f, r: %.4f, f: %.4f"
                % (dev_cost, speed, acc, p, r, f))
        else:
            current_score = acc
            logging.info("Dev: time: %.2fs speed: %.2fst/s; acc: %.4f" %
                         (dev_cost, speed, acc))

        if current_score > best_dev:
            if data.seg:
                logging.info("Exceed previous best f score: %s" % best_dev)
            else:
                logging.info("Exceed previous best acc score:%s" % best_dev)
            model_name = data.model_dir + '.' + str(idx) + ".model"
            logging.info("Save current best model in file: %s" % model_name)
            torch.save(model.state_dict(), model_name)
            best_dev = current_score
        # decode test
        speed, acc, p, r, f, _, _ = evaluate(data, model, "test")
        test_finish = time.time()
        test_cost = test_finish - dev_finish
        if data.seg:
            logging.info(
                "Test: time: %.2fs, speed: %.2fst/s; acc: %.4f, p: %.4f, r: %.4f, f: %.4f"
                % (test_cost, speed, acc, p, r, f))
        else:
            logging.info("Test: time: %.2fs, speed: %.2fst/s; acc: %.4f" %
                         (test_cost, speed, acc))
        gc.collect()

def load_model_decode(data, name):
    logging.info("Load Model from file: %s" % data.load_model_dir)
    if data.sentence_classification:
        model = SentClassifier(data)
    else:
        model = SeqLabel(data)
    # model = SeqModel(data)
    # load model need consider if the model trained in GPU and load in CPU, or vice versa
    # if not gpu:
    #     model.load_state_dict(torch.load(model_dir))
    #     # model.load_state_dict(torch.load(model_dir), map_location=lambda storage, loc: storage)
    #     # model = torch.load(model_dir, map_location=lambda storage, loc: storage)
    # else:
    #     model.load_state_dict(torch.load(model_dir))
    #     # model = torch.load(model_dir)
    model.load_state_dict(torch.load(data.load_model_dir))

    logging.info("Decode %s data, nbest: %s ..." % (name, data.nbest))
    start_time = time.time()
    speed, acc, p, r, f, pred_results, pred_scores = evaluate(
        data, model, name, data.nbest)
    end_time = time.time()
    time_cost = end_time - start_time
    if data.seg:
        logging.info(
            "%s: time:%.2fs, speed:%.2fst/s; acc: %.4f, p: %.4f, r: %.4f, f: %.4f"
            % (name, time_cost, speed, acc, p, r, f))
    else:
        logging.info("%s: time:%.2fs, speed:%.2fst/s; acc: %.4f" %
                     (name, time_cost, speed, acc))
    return pred_results, pred_scores

if __name__ == '__main__':
    utils.configure_logging()
    parser = argparse.ArgumentParser(description='Tuning with NCRF++')
    parser.add_argument(
        '--status',
        choices=['train', 'decode'],
        help='update algorithm',
        default='train')
    parser.add_argument('--config', help='Configuration File')
    args = parser.parse_args()
    data = Data()
    data.HP_gpu = torch.cuda.is_available()
    data.read_config(args.config)
    status = args.status
    logging.info("Seed num:%s" % seed_num)

    used_feature_names = list([])
    if data.feat_config is not None:
        used_feature_names = list(data.feat_config.keys())

    if status == 'train':
        logging.info("MODEL: train")
        data_initialization(data)
        data.generate_instance('train', used_feature_names)
        data.generate_instance('dev', used_feature_names)
        data.generate_instance('test', used_feature_names)
        data.build_pretrain_emb()
        train(data)
    elif status == 'decode':
        logging.info("MODEL: decode")
        data.load(data.dset_dir)
        data.read_config(args.config)
        logging.info(data.raw_dir)
        data.show_data_summary()
        data.generate_instance('raw', used_feature_names)
        logging.info("nbest: %s" % data.nbest)
        decode_results, pred_scores = load_model_decode(data, 'raw')
        if data.nbest and not data.sentence_classification:
            data.write_nbest_decoded_results(decode_results, pred_scores,
                                             'raw')
        else:
            data.write_decoded_results(decode_results, 'raw')
    else:
        logging.info(
            "Invalid argument! Please use valid arguments! (train/test/decode)"
        )

max_stars_repo_name: mahatmaWM/NCRFpp

max_issues_repo_name: mahatmaWM/NCRFpp

max_forks_repo_name: mahatmaWM/NCRFpp


text:
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  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.
""" Finetuning the library models for sequence classification on HANS."""

import logging
import os
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import numpy as np
import torch

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import is_main_process
from utils_hans import HansDataset, InputFeatures, hans_processors, hans_tasks_num_labels

logger = logging.getLogger(__name__)

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )

@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    task_name: str = field(
        metadata={"help": "The name of the task to train selected in the list: " + ", ".join(hans_processors.keys())}
    )
    data_dir: str = field(
        metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."}
    )
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )

def hans_data_collator(features: List[InputFeatures]) -> Dict[str, torch.Tensor]:
    """
    Data collator that removes the "pairID" key if present.
    """
    batch = default_data_collator(features)
    _ = batch.pop("pairID", None)
    return batch

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
        )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        training_args.local_rank,
        training_args.device,
        training_args.n_gpu,
        bool(training_args.local_rank != -1),
        training_args.fp16,
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed
    set_seed(training_args.seed)

    try:
        num_labels = hans_tasks_num_labels[data_args.task_name]
    except KeyError:
        raise ValueError("Task not found: %s" % (data_args.task_name))

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=data_args.task_name,
        cache_dir=model_args.cache_dir,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
    )

    # Get datasets
    train_dataset = (
        HansDataset(
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            task=data_args.task_name,
            max_seq_length=data_args.max_seq_length,
            overwrite_cache=data_args.overwrite_cache,
        )
        if training_args.do_train
        else None
    )
    eval_dataset = (
        HansDataset(
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            task=data_args.task_name,
            max_seq_length=data_args.max_seq_length,
            overwrite_cache=data_args.overwrite_cache,
            evaluate=True,
        )
        if training_args.do_eval
        else None
    )

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=hans_data_collator,
    )

    # Training
    if training_args.do_train:
        trainer.train(
            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
        )
        trainer.save_model()
        # For convenience, we also re-save the tokenizer to the same directory,
        # so that you can share your model easily on huggingface.co/models =)
        if trainer.is_world_master():
            tokenizer.save_pretrained(training_args.output_dir)

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        output = trainer.predict(eval_dataset)
        preds = output.predictions
        preds = np.argmax(preds, axis=1)

        mismatch_eval_file = os.path.join(training_args.output_dir, "hans_mismatched.txt")

        pair_ids = [ex.pairID for ex in eval_dataset]
        output_eval_file = os.path.join(training_args.output_dir, "hans_predictions.txt")
        label_list = eval_dataset.get_labels()
        if trainer.is_world_master():
            with open(output_eval_file, "w") as writer:
                with open(mismatch_eval_file, "w") as mismatch_writer:
                    writer.write("pairID,gold_label\n")
                    for pid, pred in zip(pair_ids, preds):
                        writer.write("ex" + str(pid) + "," + label_list[int(pred)] + "\n")

                        if int(pred) != int(eval_dataset[pid].label):
                            mismatch_writer.write(str(pid) + "\n")

        trainer._log(output.metrics)

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()

if __name__ == "__main__":
    main()

ext: py

max_stars_repo_name: kiromaru/transformers

max_issues_repo_name: kiromaru/transformers

max_forks_repo_name: kiromaru/transformers

reason: import numpy


text:
from pathlib import Path

import numpy as np
import soundfile

from espnet2.fileio.sound_scp import SoundScpReader

def test_SoundScpReader(tmp_path: Path):
    audio_path1 = tmp_path / "a1.wav"
    audio1 = np.random.randint(-100, 100, 16, dtype=np.int16)
    audio_path2 = tmp_path / "a2.wav"
    audio2 = np.random.randint(-100, 100, 16, dtype=np.int16)

    soundfile.write(audio_path1, audio1, 16)
    soundfile.write(audio_path2, audio2, 16)

    p = tmp_path / "dummy.scp"
    with p.open("w") as f:
        f.write(f"abc {audio_path1}\n")
        f.write(f"def {audio_path2}\n")

    desired = {"abc": (16, audio1), "def": (16, audio2)}
    target = SoundScpReader(p, normalize=False, dtype=np.int16)

    for k in desired:
        rate1, t = target[k]
        rate2, d = desired[k]
        assert rate1 == rate2
        np.testing.assert_array_equal(t, d)

    assert len(target) == len(desired)
    assert "abc" in target
    assert "def" in target
    assert tuple(target.keys()) == tuple(desired)
    assert tuple(target) == tuple(desired)
    assert target.get_path("abc") == str(audio_path1)
    assert target.get_path("def") == str(audio_path2)

def test_SoundScpReader_normalize(tmp_path: Path):
    audio_path1 = tmp_path / "a1.wav"
    audio1 = np.random.randint(-100, 100, 16, dtype=np.int16)
    audio_path2 = tmp_path / "a2.wav"
    audio2 = np.random.randint(-100, 100, 16, dtype=np.int16)

    audio1 = audio1.astype(np.float64) / (np.iinfo(np.int16).max + 1)
    audio2 = audio2.astype(np.float64) / (np.iinfo(np.int16).max + 1)

    soundfile.write(audio_path1, audio1, 16)
    soundfile.write(audio_path2, audio2, 16)

    p = tmp_path / "dummy.scp"
    with p.open("w") as f:
        f.write(f"abc {audio_path1}\n")
        f.write(f"def {audio_path2}\n")

    desired = {"abc": (16, audio1), "def": (16, audio2)}
    target = SoundScpReader(p, normalize=True)

    for k in desired:
        rate1, t = target[k]
        rate2, d = desired[k]
        assert rate1 == rate2
        np.testing.assert_array_equal(t, d)


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

# Copyright 2018 IBM.
#
# 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.
# =============================================================================

from functools import partial

import numpy as np

from qiskit_aqua import QuantumAlgorithm, get_optimizer_instance
from qiskit_aqua.svm import (eval_cost_function, eval_cost_function_with_unlabeled_data,
                             set_print_info, entangler_map_creator)

class SVM_Variational(QuantumAlgorithm):

    SVM_VARIATIONAL_CONFIGURATION = {
        'name': 'SVM_Variational',
        'description': 'SVM_Variational Algorithm',
        'input_schema': {
            '$schema': 'http://json-schema.org/schema#',
            'id': 'SVM_Variational_schema',
            'type': 'object',
            'properties': {

                'circuit_depth': {
                    'type': 'integer',
                    'default': 3,
                    'minimum': 3
                },
                'print_info': {
                    'type': 'boolean',
                    'default': False
                }
            },
            'additionalProperties': False
        },
        'depends': ['optimizer'],
        'problems': ['svm_classification'],
        'defaults': {
            'optimizer': {
                'name': 'SPSA'
            }
        }
    }

    def __init__(self, configuration=None):
        super().__init__(configuration or self.SVM_VARIATIONAL_CONFIGURATION.copy())
        self._ret = {}

    def init_params(self, params, algo_input):
        SVMQK_params = params.get(QuantumAlgorithm.SECTION_KEY_ALGORITHM)

        circuit_depth = SVMQK_params.get('circuit_depth')
        print_info = SVMQK_params.get('print_info')

        # Set up optimizer
        opt_params = params.get(QuantumAlgorithm.SECTION_KEY_OPTIMIZER)
        optimizer = get_optimizer_instance(opt_params['name'])
        # hard-coded params if SPSA is used.
        if opt_params['name'] == 'SPSA' and opt_params['parameters'] is None:
            opt_params['parameters'] = np.asarray([4.0, 0.1, 0.602, 0.101, 0.0])
        optimizer.init_params(opt_params)
        optimizer.set_options(save_steps=10)

        self.init_args(algo_input.training_dataset, algo_input.test_dataset, algo_input.datapoints,
                        optimizer, circuit_depth, print_info)

    def init_args(self, training_dataset, test_dataset, datapoints, optimizer, circuit_depth=3, print_info=False):
        if 'statevector' in self._backend:
            raise ValueError('Selected backend  "{}" does not support measurements.'.format(self._backend))

        self.training_dataset = training_dataset
        self.test_dataset = test_dataset
        self.datapoints = datapoints
        self.class_labels = list(self.training_dataset.keys())

        self.num_of_qubits = self.auto_detect_qubitnum(training_dataset) # auto-detect mode
        self.entangler_map = entangler_map_creator(self.num_of_qubits)
        self.coupling_map = None  # the coupling_maps gates allowed on the device
        self.initial_layout = None
        self.shots = self._execute_config['shots']
        self.backend = self._backend
        self.circuit_depth = circuit_depth
        self.print_info = print_info
        self.optimizer = optimizer

        set_print_info(print_info)

    def auto_detect_qubitnum(self, training_dataset):
        auto_detected_size = -1
        for key in training_dataset:
            val = training_dataset[key]
            for item in val:
                auto_detected_size = len(item)
                return auto_detected_size
        return auto_detected_size

    def train(self, training_input, class_labels):
        initial_theta = self.random.randn(2 * self.num_of_qubits * (self.circuit_depth + 1))

        eval_cost_function_partial = partial(eval_cost_function, self.entangler_map, self.coupling_map,
                                             self.initial_layout, self.num_of_qubits, self.circuit_depth,
                                             training_input, class_labels, self.backend, self.shots, self._random_seed, True)

        def objective_function(theta):
            return eval_cost_function_partial(theta)[0]

        theta_best, cost_final, _ = self.optimizer.optimize(
            initial_theta.shape[0], objective_function, initial_point=initial_theta)
        # costs = cost_final  # , cost_plus, cost_minus

        self._ret['opt_params'] = theta_best
        self._ret['training_loss'] = cost_final

    def test(self, test_input, class_labels):
        total_cost, std_cost, success_ratio, predicted_labels = \
            eval_cost_function(self.entangler_map, self.coupling_map, self.initial_layout, self.num_of_qubits,
                               self.circuit_depth, test_input, class_labels, self.backend, self.shots, self._random_seed,
                               train=False, theta=self._ret['opt_params'])

        if self.print_info:
            print('Classification success for this set is  %s %%  \n' % (100.0 * success_ratio))
        return success_ratio

    def predict(self, input_datapoints, class_labels):
        predicted_labels = eval_cost_function_with_unlabeled_data(self.entangler_map, self.coupling_map,
                                                                  self.initial_layout, self.num_of_qubits,
                                                                  self.circuit_depth, input_datapoints, class_labels,
                                                                  self.backend, self.shots, self._random_seed,
                                                                  train=False, theta=self._ret['opt_params'])
        return predicted_labels

    def run(self):
        if self.training_dataset is None:
            self._ret['error'] = 'training dataset is missing! please provide it'
            return self._ret

        num_of_qubits = self.auto_detect_qubitnum(self.training_dataset) # auto-detect mode
        if num_of_qubits == -1:
            self._ret['error'] = 'Something wrong with the auto-detection of num_of_qubits'
            return self._ret
        if num_of_qubits != 2 and num_of_qubits != 3:
            self._ret['error'] = 'You should lower the feature size to 2 or 3 using PCA first!'
            return self._ret

        self.train(self.training_dataset, self.class_labels)

        if self.test_dataset is not None:
            success_ratio = self.test(self.test_dataset, self.class_labels)
            self._ret['test_success_ratio'] = success_ratio

        if self.datapoints is not None:
            predicted_labels = self.predict(self.datapoints, self.class_labels)
            self._ret['predicted_labels'] = predicted_labels

        return self._ret

max_stars_repo_name: NunoEdgarGFlowHub/aqua

max_issues_repo_name: NunoEdgarGFlowHub/aqua

max_forks_repo_name: NunoEdgarGFlowHub/aqua


text:
import numpy as np
import time
import gym
# import mujoco_py as mjc
import pdb

from mopo.parallel.parallel_hopper import ParallelHopperEnv

N = 10000
horizon = 1

env = gym.make('Hopper-v2')
parallel = ParallelHopperEnv(N)

env.reset()
t0_p = time.time()
obs = parallel.reset()
vec_qpos_init, vec_qvel_init = parallel.vec_qpos_qvel()
obs_p = [obs]
act_p = []
for i in range(horizon):
	act = [parallel.action_space.sample() for _ in range(parallel.N)]
	obs, rew, term, info = parallel.step(act)
	obs_p.append(obs)
	act_p.append(act)
obs_p = np.array(obs_p)
act_p = np.array(act_p)
t1_p = time.time()

##
t0_s = time.time()
obs_s = []
for sim in range(parallel.N):
	print(sim)
	qpos = vec_qpos_init[sim]
	qvel = vec_qvel_init[sim]
	actions = act_p[:,sim]

	env.set_state(qpos, qvel)
	obs = env.env._get_obs()
	obs_s.append([obs])
	for act in actions:
		obs, rew, term, info = env.step(act)
		obs_s[-1].append(obs)
t1_s = time.time()

## [ num_sims x horizon x obs_dim ] --> [ horizon x num_sims x obs_dim ]
obs_s = np.array(obs_s).transpose(1,0,2)
diff = np.abs(obs_p - obs_s).max()
print(obs_p.shape)

time_p = t1_p - t0_p
time_s = t1_s - t0_s

print('environments: {} | horizon: {} | diff: '.format(N, horizon, diff))
print('parallel: {} | sequential: {}'.format(time_p, time_s))
pdb.set_trace()


text:
# Copyright (c) Microsoft Corporation
# Licensed under the MIT License.

from typing import Callable, Dict, List

import numpy as np
import pandas as pd

from responsibleai.databalanceanalysis.balance_measures import BalanceMeasures
from responsibleai.databalanceanalysis.constants import Measures

def _get_generalized_entropy_index(
    benefits: np.array,
    alpha: float,
    use_abs_val: bool,
    error_tolerance: float = 1e-12,
) -> float:
    """
    Returns the general entropy index, a measure of redundancy in the data.
    https://en.wikipedia.org/wiki/Generalized_entropy_index

    :param benefits: The probabilities of each feature value combination
        occurring in the data.
    :type benefits: np.array
    :param alpha: A parameter which regulates the weight given to distances
        between benefits at different parts of the distribution.
    :type alpha: float
    :param use_abs_val: Whether to use the absolute value of the benefits
        or not.
    :type use_abs_val: bool
    :param error_tolerance: The error tolerance to use when calculating
        the generalized entropy index.
    :type error_tolerance: float
    :return: The generalized entropy index.
    :rtype: float
    """
    if use_abs_val:
        benefits = np.absolute(benefits)

    benefits_mean = np.mean(benefits)
    norm_benefits = benefits / benefits_mean
    count = norm_benefits.size

    if abs(alpha - 1.0) < error_tolerance:
        gei = np.sum(norm_benefits * np.log(norm_benefits)) / count
    elif abs(alpha) < error_tolerance:
        gei = np.sum(-np.log(norm_benefits)) / count
    else:
        gei = np.sum(np.power(norm_benefits, alpha) - 1.0) / (
            count * alpha * (alpha - 1.0)
        )

    return gei

def get_atkinson_index(
    benefits: np.array,
    epsilon: float = 1.0,
    error_tolerance: float = 1e-12,
) -> float:
    """
    Returns the Atkinson Index, a measure of income inequality that
    indicates the percentage of the data that needs to be foregone
    to have equal shares of income.
    https://en.wikipedia.org/wiki/Atkinson_index

    :param benefits: The probabilities of each feature value combination
        occurring in the data.
    :type benefits: np.array
    :param epsilon: A parameter that regulates the weight given
        to the distances between benefits.
    :type epsilon: float
    :param error_tolerance: The error tolerance to use when calculating
        the Atkinson Index.
    :type error_tolerance: float
    :return: The Atkinson Index.
    :rtype: float
    """
    count = benefits.size
    benefits_mean = np.mean(benefits)
    norm_benefits = benefits / benefits_mean
    alpha = 1 - epsilon

    if abs(alpha) < error_tolerance:
        ati = 1.0 - np.power(np.prod(norm_benefits), 1.0 / count)
    else:
        power_mean = np.sum(np.power(norm_benefits, alpha)) / count
        ati = 1.0 - np.power(power_mean, 1.0 / alpha)

    return ati

def get_theil_t_index(benefits: np.array) -> float:
    """
    Returns Theil T Index, a measure of income inequality that is more
    sensitive to differences at the top of the distribution.
    https://en.wikipedia.org/wiki/Theil_index

    :param benefits: The probabilities of each feature value combination
        occurring in the data.
    :type benefits: np.array
    :return: The Theil T Index.
    :rtype: float
    """
    return _get_generalized_entropy_index(
        benefits=benefits, alpha=1.0, use_abs_val=True
    )

def get_theil_l_index(benefits: np.array) -> float:
    """
    Returns Theil L Index, a measure of income inequality that is
    more sensitive to differences at the lower end of the distribution.
    https://en.wikipedia.org/wiki/Theil_index

    :param benefits: The probabilities of each feature value combination
        occurring in the data.
    :type benefits: np.array
    :return: The Theil L Index.
    :rtype: float
    """
    return _get_generalized_entropy_index(
        benefits=benefits, alpha=0.0, use_abs_val=True
    )

class AggregateBalanceMeasures(BalanceMeasures):
    """
    This class computes a set of aggregated balance measures that represents
    how balanced the given dataset is along the given columns of interest.
    """

    AGGREGATE_METRICS: Dict[Measures, Callable[[np.array], float]] = {
        Measures.THEIL_L_INDEX: get_theil_l_index,
        Measures.THEIL_T_INDEX: get_theil_t_index,
        Measures.ATKINSON_INDEX: get_atkinson_index,
    }

    def __init__(self, cols_of_interest: List[str]) -> None:
        """
        Construct a AggregateBalanceMeasures class for calculating
        aggregate balance measures on a dataset.

        :param cols_of_interest: The list of columns to compute
            aggregate balance measures on.
        :type cols_of_interest: List[str]
        """
        super().__init__(cols_of_interest=cols_of_interest)

    def measures(self, dataset: pd.DataFrame) -> pd.DataFrame:
        """
        Returns aggregate balance measures for the given dataset and
        columns of interest.

        The following measures are computed:
        * Atkinson Index - https://en.wikipedia.org/wiki/Atkinson_index
        * Theil Index (L and T) - https://en.wikipedia.org/wiki/Theil_index

        :param dataset: The dataset to compute aggregate measures on.
        :type dataset: pd.DataFrame
        :return: A pandas DataFrame that contains the names and values of
            aggregate balance measures for the given dataset and
            columns of interest.
        :rtype: pd.DataFrame
        """
        aggregate_measures_dict = {}
        benefits = (
            dataset.groupby(self.cols_of_interest).size() / dataset.shape[0]
        )

        for measure, func in self.AGGREGATE_METRICS.items():
            aggregate_measures_dict[measure.value] = [func(benefits)]

        aggregate_measures = pd.DataFrame.from_dict(aggregate_measures_dict)
        return aggregate_measures

max_stars_repo_name: imatiach-msft/responsible-ai-toolbox

max_issues_repo_name: imatiach-msft/responsible-ai-toolbox

max_forks_repo_name: imatiach-msft/responsible-ai-toolbox


text:
from this import d
import matplotlib.dates
import numpy as np
from floodsystem.station import MonitoringStation

def polyfit(dates, levels, p):
    x = matplotlib.dates.date2num(dates)
    x0 = x[0]
    y = levels
    p_coeff = np.polyfit(x - x0, y, p)
    poly = np.poly1d(p_coeff)
    return poly, x0

def floodrisk(station, dates, levels):
    risk = 0
    #determine risk with relative water level
    global relative_level
    if station.typical_range_consistent() == False:
        pass
    elif station.latest_level == None:
        pass
    else:
        relative_level = relative_level = (station.latest_level - station.typical_range[0])/(station.typical_range[1]-station.typical_range[0])

    if relative_level > 1.5:
        risk += 5
    elif relative_level > 1:
        risk += 3
    elif relative_level > 0.5:
        risk += 1
    else:
        risk += 0

    #determine risk with tendency of change of water level
    global gradient
    x = matplotlib.dates.date2num(dates)
    if len(x) == 0:
        pass
    else:
        x1 = np.linspace(x[0], x[-1], 100)
        if len(x - x[0]) != len(levels):
            pass
        elif len(x - x[0]) != len(levels):
            pass
        else:
            try:
                p_coeff = np.polyfit(x - x[0], levels, 4)
                poly = np.poly1d(p_coeff)
                gradient = (poly(x1[-1] - x[0]) - poly(x1[0] - x[0]))
            except TypeError:
                pass

    if gradient > 1:
        risk += 3
    elif gradient > 0.5:
        risk += 1
    elif gradient < 0:
        risk -= 1
    else:
        risk += 0

    if risk <= 1:
        return 'low'
    elif risk == 2:
        return 'moderate'
    elif risk == 3:
        return 'high'
    else:
        return 'severe'


text:
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
r"""
Basic training script for PyTorch
"""

# Set up custom environment before nearly anything else is imported
# NOTE: this should be the first import (no not reorder)
from maskrcnn_benchmark.utils.env import setup_environment  # noqa F401 isort:skip

import argparse, logging, time, datetime
import os
import numpy as np

import torch
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.data import make_data_loader
from maskrcnn_benchmark.solver import make_lr_scheduler
from maskrcnn_benchmark.solver import make_optimizer
from maskrcnn_benchmark.engine.inference import inference
# from maskrcnn_benchmark.engine.trainer import do_train
from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.utils.collect_env import collect_env_info
from maskrcnn_benchmark.utils.comm import synchronize, \
    get_rank, is_pytorch_1_1_0_or_later
from maskrcnn_benchmark.utils.imports import import_file
from maskrcnn_benchmark.utils.logger import setup_logger
from maskrcnn_benchmark.utils.miscellaneous import mkdir

from maskrcnn_benchmark.utils.comm import get_world_size, is_pytorch_1_1_0_or_later
from maskrcnn_benchmark.utils.metric_logger import MetricLogger

def reduce_loss_dict(loss_dict):
    """
    Reduce the loss dictionary from all processes so that process with rank
    0 has the averaged results. Returns a dict with the same fields as
    loss_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return loss_dict
    with torch.no_grad():
        loss_names = []
        all_losses = []
        for k in sorted(loss_dict.keys()):
            loss_names.append(k)
            all_losses.append(loss_dict[k])
        all_losses = torch.stack(all_losses, dim=0)
        torch.distributed.reduce(all_losses, dst=0)
        if torch.distributed.get_rank() == 0:
            # only main process gets accumulated, so only divide by
            # world_size in this case
            all_losses /= world_size
        reduced_losses = {k: v for k, v in zip(loss_names, all_losses)}
    return reduced_losses

def do_train(
    model,
    data_loader,
    optimizer,
    scheduler,
    checkpointer,
    device,
    checkpoint_period,
    arguments,
):
    logger = logging.getLogger("maskrcnn_benchmark.trainer")
    logger.info("Start training")
    meters = MetricLogger(delimiter="  ")
    max_iter = len(data_loader)
    start_iter = arguments["iteration"]
    model.train()
    start_training_time = time.time()
    end = time.time()
    pytorch_1_1_0_or_later = is_pytorch_1_1_0_or_later()
    for iteration, (images, targets, _) in enumerate(data_loader, start_iter):
        data_time = time.time() - end
        iteration = iteration + 1
        arguments["iteration"] = iteration

        # in pytorch >= 1.1.0, scheduler.step() should be run after optimizer.step()
        if not pytorch_1_1_0_or_later:
            scheduler.step()

        images = images.to(device)
        targets = [target.to(device) for target in targets]

        loss_dict = model(images, targets)

        losses = sum(loss for loss in loss_dict.values())
        # print('iteration:{}, losses:{}'.format(iteration, losses))
        if losses > 1e5 or torch.isnan(losses):
            import pdb; pdb.set_trace()

        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = reduce_loss_dict(loss_dict)
        losses_reduced = sum(loss for loss in loss_dict_reduced.values())
        meters.update(loss=losses_reduced, **loss_dict_reduced)

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        if pytorch_1_1_0_or_later:
            scheduler.step()

        batch_time = time.time() - end
        end = time.time()
        meters.update(time=batch_time, data=data_time)

        eta_seconds = meters.time.global_avg * (max_iter - iteration)
        eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))

        if iteration % 20 == 0 or iteration == max_iter:
            logger.info(
                meters.delimiter.join(
                    [
                        "eta: {eta}",
                        "iter: {iter}",
                        "{meters}",
                        "lr: {lr:.6f}",
                        "max mem: {memory:.0f}",
                    ]
                ).format(
                    eta=eta_string,
                    iter=iteration,
                    meters=str(meters),
                    lr=optimizer.param_groups[0]["lr"],
                    memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0,
                )
            )
        if iteration % checkpoint_period == 0:
            checkpointer.save("model_{:07d}".format(iteration), **arguments)
            # import pdb; pdb.set_trace()
            run_test(cfg, model, distributed=False, test_epoch=iteration)
            model.train()
        if iteration == max_iter:
            checkpointer.save("model_final", **arguments)

    total_training_time = time.time() - start_training_time
    total_time_str = str(datetime.timedelta(seconds=total_training_time))
    logger.info(
        "Total training time: {} ({:.4f} s / it)".format(
            total_time_str, total_training_time / (max_iter)
        )
    )

from torch import nn
import torch.nn.functional as F
from maskrcnn_benchmark.modeling.backbone import resnet
from maskrcnn_benchmark.modeling.backbone import fpn as fpn_module
from maskrcnn_benchmark.modeling.make_layers import conv_with_kaiming_uniform
from maskrcnn_benchmark.modeling.rpn.fcos.fcos import FCOSModule
from maskrcnn_benchmark.structures.bounding_box import BoxList

class MaskFirst(nn.Module):
    def __init__(self, cfg):
        super(MaskFirst, self).__init__()
        self.cfg = cfg
        self.r50 = resnet.ResNet(cfg)
        in_channels_stage2 = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS
        out_channels = cfg.MODEL.RESNETS.BACKBONE_OUT_CHANNELS
        in_channels_p6p7 = in_channels_stage2 * 8 if cfg.MODEL.RETINANET.USE_C5 \
            else out_channels
        self.fpn = fpn_module.FPN(
            in_channels_list=[
                0,
                in_channels_stage2 * 2,
                in_channels_stage2 * 4,
                in_channels_stage2 * 8,
            ],
            out_channels=out_channels,
            conv_block=conv_with_kaiming_uniform(
                cfg.MODEL.FPN.USE_GN, cfg.MODEL.FPN.USE_RELU
            ),
            top_blocks=fpn_module.LastLevelP6P7(in_channels_p6p7, out_channels),
        )
        # self.mask_cls = nn.Conv2d(256, 80, 1)
        self.loss_evaluators = []
        self.fcos = FCOSModule(cfg, 256)
        self.p7p6conv = nn.Conv2d(256, 1, 1)
        # per_proposal_mask 应适当多层 3～5层3*3
        # self.mask_init_conv = nn.Conv2d(256+1, 1, 3, padding=1)
        self.mask_init_conv = nn.Conv2d(256+1, 1, 7, padding=3)
        self.score = nn.Conv2d(256, 1, 1)
        # self.criterion = nn.CrossEntropyLoss()
        # 这里先用L1 loss 后续测试l2 loss 以及cross entropy

    def forward(self, image, targets=None):
        x_img = image.tensors
        xs_r50 = self.r50(x_img)
        fs_fpn = self.fpn(xs_r50)
        fs_fpn_p7 = fs_fpn[4]
        # 考虑改用mask score rcnn
        x_score = self.score(fs_fpn_p7)
        # x_fcos = self.fcos(image, fs_fpn, targets)
        # proposals, proposal_losses = self.rpn(images, features, targets)
        # print(type(image), type(targets))
        # 先不做nms， bbox格式 x, y, x_shape, y_shape
        N, C, x_shape, y_shape = fs_fpn_p7.shape
        fake_bbox = torch.tensor([[x,y,x_shape,y_shape] for x in range(x_shape) for y in range(y_shape)])
        fake_bbox = torch.tensor(fake_bbox)

        results = []
        masks_p7 = []
        losses = []
        for i in range(N):
            h_ori, w_ori = image.image_sizes[i]
            boxlist = BoxList(fake_bbox, (int(w_ori), int(h_ori)), mode="xyxy")
            # boxlist.add_field("labels", per_class)
            boxlist.add_field("scores", x_score)
            # import pdb; pdb.set_trace()
            target_levels = self._init_target(tuple(x_img.shape[-2:]), x_img.device, targets[i])
            init_inst_pyramids = self.init_inst_pyramid(fs_fpn_p7[i])
            loss7 = self.loss_of_pyramid_7(init_inst_pyramids, target_levels)
            # losses.append(torch.tensor(loss7, device=x_img.device).sum())
            losses.append(sum(loss for loss in loss7))
        # print([t.shape for t in xs_r50])
        # import pdb; pdb.set_trace()
        # losses = torch.tensor(losses, device=x_img.device).sum()
        losses = sum(loss for loss in losses)
        return {'loss_mask_7': losses}

    def loss_of_pyramid_7(self, pyramids, target_levels):
        # import pdb; pdb.set_trace()
        losses = []
        for pyramid in pyramids:
            mask_7 = torch.sigmoid(pyramid.masks[7])
            target_7 = target_levels[7]
            target_pos = pyramid.pos
            target_idx = target_7[0, :, target_pos[0], target_pos[1]].nonzero()
            # TODO: target的两类重复问题：1. 两个object在同一个pixel上；
            #   2. 一个object占用多个pixel导致单个object有两个pyramid，需要merge，或nms
            if len(target_idx):
                if len(target_idx) > 1:
                    # print('xxxxxxxxxxx\n:', target_7.shape, target_idx)
                    # import pdb; pdb.set_trace()
                    # areas = [target_7[0,idx, ...].sum() for idx in target_idx]
                    target_idx = target_7[0, target_idx].view(len(target_idx),-1).sum(1).argmin().unsqueeze(0).unsqueeze(0)
                target_mask = target_7[:, [target_idx]]
                # target_mask = target_7[:, target_idx]
                pyramid.bind_target(target_idx)
                # loss = (mask_7 - target_mask.float()).abs().sum()
                loss = (mask_7 - target_mask.float()).abs().mean()
            else:   # 没有对应target，说明对应unknown类
                # loss = (target_7.float().sum(dim=1).unsqueeze(1) * mask_7).sum()
                loss = (target_7.float().sum(dim=1).unsqueeze(1) * mask_7).mean()
            losses.append(loss)
        return losses

    def init_inst_pyramid(self, fs_fpn_p7_i):
        # import pdb; pdb.set_trace()
        _, h, w = fs_fpn_p7_i.shape
        insts_per_img = []
        for h_i in range(h):
            for w_i in range(w):
                inst_pyra = InstancePyramid(h_i, w_i, 7, (h,w))
                inst_pyra._init_mask(fs_fpn_p7_i, self.mask_init_conv)
                insts_per_img.append(inst_pyra)

        return insts_per_img

    def _init_target(self, img_tensor_shape, device, target=None):
        target_ori_mask = target.get_field('masks').get_mask_tensor().unsqueeze(0).to(device)

        # import pdb; pdb.set_trace()
        target_shape = (1, target_ori_mask.shape[-3]) + img_tensor_shape
        target_mask_pad_to_img = target_ori_mask.new(*target_shape).zero_()
        target_mask_pad_to_img[:,:,:target.size[1], :target.size[0]] = target_ori_mask

        target_levels = {}
        target_levels[0] = target_mask_pad_to_img
        level_shape = ((img_tensor_shape[0]+1)//2, (img_tensor_shape[1]+1)//2)
        target_levels[1] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)
        level_shape = ((level_shape[0]+1)//2, (level_shape[1]+1)//2)
        target_levels[2] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)
        level_shape = ((level_shape[0]+1)//2, (level_shape[1]+1)//2)
        target_levels[3] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)
        level_shape = ((level_shape[0]+1)//2, (level_shape[1]+1)//2)
        target_levels[4] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)
        level_shape = ((level_shape[0]+1)//2, (level_shape[1]+1)//2)
        target_levels[5] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)
        level_shape = ((level_shape[0]+1)//2, (level_shape[1]+1)//2)
        target_levels[6] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)
        level_shape = ((level_shape[0]+1)//2, (level_shape[1]+1)//2)
        target_levels[7] = F.interpolate(target_mask_pad_to_img.float(), level_shape, mode='nearest').type(target_mask_pad_to_img.dtype)

        # import pdb; pdb.set_trace()
        # print([t.shape for t in target_levels.values()])
        return target_levels

    def forward_for_single_feature_map(
            self, locations, box_cls,
            box_regression, centerness,
            image_sizes):
        """
        Arguments:
            anchors: list[BoxList]
            box_cls: tensor of size N, A * C, H, W
            box_regression: tensor of size N, A * 4, H, W
        """
        N, C, H, W = box_cls.shape

        # put in the same format as locations
        box_cls = box_cls.view(N, C, H, W).permute(0, 2, 3, 1)
        box_cls = box_cls.reshape(N, -1, self.num_classes - 1).sigmoid()
        box_regression = box_regression.view(N, self.dense_points * 4, H, W).permute(0, 2, 3, 1)
        box_regression = box_regression.reshape(N, -1, 4)
        centerness = centerness.view(N, self.dense_points, H, W).permute(0, 2, 3, 1)
        centerness = centerness.reshape(N, -1).sigmoid()

        candidate_inds = box_cls > self.pre_nms_thresh
        pre_nms_top_n = candidate_inds.view(N, -1).sum(1)
        pre_nms_top_n = pre_nms_top_n.clamp(max=self.pre_nms_top_n)

        # multiply the classification scores with centerness scores
        box_cls = box_cls * centerness[:, :, None]

        results = []
        for i in range(N):
            per_box_cls = box_cls[i]
            per_candidate_inds = candidate_inds[i]
            per_box_cls = per_box_cls[per_candidate_inds]

            per_candidate_nonzeros = per_candidate_inds.nonzero()
            per_box_loc = per_candidate_nonzeros[:, 0]
            per_class = per_candidate_nonzeros[:, 1] + 1

            per_box_regression = box_regression[i]
            per_box_regression = per_box_regression[per_box_loc]
            per_locations = locations[per_box_loc]

            per_pre_nms_top_n = pre_nms_top_n[i]

            if per_candidate_inds.sum().item() > per_pre_nms_top_n.item():
                per_box_cls, top_k_indices = \
                    per_box_cls.topk(per_pre_nms_top_n, sorted=False)
                per_class = per_class[top_k_indices]
                per_box_regression = per_box_regression[top_k_indices]
                per_locations = per_locations[top_k_indices]

            detections = torch.stack([
                per_locations[:, 0] - per_box_regression[:, 0],
                per_locations[:, 1] - per_box_regression[:, 1],
                per_locations[:, 0] + per_box_regression[:, 2],
                per_locations[:, 1] + per_box_regression[:, 3],
            ], dim=1)

            h, w = image_sizes[i]
            boxlist = BoxList(detections, (int(w), int(h)), mode="xyxy")
            boxlist.add_field("labels", per_class)
            boxlist.add_field("scores", per_box_cls)
            boxlist = boxlist.clip_to_image(remove_empty=False)
            boxlist = remove_small_boxes(boxlist, self.min_size)
            results.append(boxlist)

        return results

class InstancePyramid():
    def __init__(self, pos_h, pos_w, init_level, level_size, target=None):
        self.init_level = init_level
        self.level_size = level_size
        self.pos_h = pos_h
        self.pos_w = pos_w
        self.pos = (self.pos_h, self.pos_w)
        self.masks = {}
        self.target_idx = None
        self.score = 1.0

    def _init_mask(self, feature, mask_init_conv):
        # import pdb; pdb.set_trace()
        f_0 = torch.zeros_like(feature[0]).unsqueeze(0)
        if feature.shape[-2:] == self.level_size:
            f_0[0, self.pos_h, self.pos_w] = 1.0
        feature_pos = torch.cat((feature, f_0)).unsqueeze(0)
        self.masks[self.init_level] = mask_init_conv(feature_pos)

    def bind_target(self, idx):
        self.target_idx = idx

    def _calcu_loss(self):
        import pdb; pdb.set_trace()

def train(cfg, local_rank, distributed):
    # model = build_detection_model(cfg)
    model = MaskFirst(cfg)
    device = torch.device(cfg.MODEL.DEVICE)
    model.to(device)

    if cfg.MODEL.USE_SYNCBN:
        assert is_pytorch_1_1_0_or_later(), \
            "SyncBatchNorm is only available in pytorch >= 1.1.0"
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

    optimizer = make_optimizer(cfg, model)
    scheduler = make_lr_scheduler(cfg, optimizer)

    if distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[local_rank], output_device=local_rank,
            # this should be removed if we update BatchNorm stats
            broadcast_buffers=False,
        )

    arguments = {}
    arguments["iteration"] = 0

    output_dir = cfg.OUTPUT_DIR

    save_to_disk = get_rank() == 0
    checkpointer = DetectronCheckpointer(
        cfg, model, optimizer, scheduler, output_dir, save_to_disk
    )
    extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
    arguments.update(extra_checkpoint_data)

    data_loader = make_data_loader(
        cfg,
        is_train=True,
        is_distributed=distributed,
        start_iter=arguments["iteration"],
    )

    checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD

    do_train(
        model,
        data_loader,
        optimizer,
        scheduler,
        checkpointer,
        device,
        checkpoint_period,
        arguments,
    )

    return model

def get_neat_inference_result(coco_eval):
    def _summarize( ap=1, iouThr=None, areaRng='all', maxDets=100 ):
        p = coco_eval.params
        # import pdb; pdb.set_trace()
        iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
        titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
        typeStr = '(AP)' if ap==1 else '(AR)'
        iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
            if iouThr is None else '{:0.2f}'.format(iouThr)

        aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
        mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
        if ap == 1:
            # dimension of precision: [TxRxKxAxM]
            # [p.iouThrs * p.recThrs * p.catIds * p.areaRng * p.maxDets]
            s = coco_eval.eval['precision']
            # IoU
            if iouThr is not None:
                t = np.where(iouThr == p.iouThrs)[0]
                s = s[t]
            s = s[:,:,:,aind,mind]
        else:
            # dimension of recall: [TxKxAxM]
            s = coco_eval.eval['recall']
            if iouThr is not None:
                t = np.where(iouThr == p.iouThrs)[0]
                s = s[t]
            s = s[:,:,aind,mind]
        if len(s[s>-1])==0:
            mean_s = -1
        else:
            mean_s = np.mean(s[s>-1])
        print(iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s))
        summaryStr = iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s)
        return mean_s, summaryStr
    def _summarizeDets():
        stats = np.zeros((12,))
        summaryStrs = [None] * 12
        stats[0], summaryStrs[0] = _summarize(1)
        stats[1], summaryStrs[1] = _summarize(1, iouThr=.5, maxDets=coco_eval.params.maxDets[2])
        stats[2], summaryStrs[2] = _summarize(1, iouThr=.75, maxDets=coco_eval.params.maxDets[2])
        stats[3], summaryStrs[3] = _summarize(1, areaRng='small', maxDets=coco_eval.params.maxDets[2])
        stats[4], summaryStrs[4] = _summarize(1, areaRng='medium', maxDets=coco_eval.params.maxDets[2])
        stats[5], summaryStrs[5] = _summarize(1, areaRng='large', maxDets=coco_eval.params.maxDets[2])
        stats[6], summaryStrs[6] = _summarize(0, maxDets=coco_eval.params.maxDets[0])
        stats[7], summaryStrs[7] = _summarize(0, maxDets=coco_eval.params.maxDets[1])
        stats[8], summaryStrs[8] = _summarize(0, maxDets=coco_eval.params.maxDets[2])
        stats[9], summaryStrs[9] = _summarize(0, areaRng='small', maxDets=coco_eval.params.maxDets[2])
        stats[10], summaryStrs[10] = _summarize(0, areaRng='medium', maxDets=coco_eval.params.maxDets[2])
        stats[11], summaryStrs[11] = _summarize(0, areaRng='large', maxDets=coco_eval.params.maxDets[2])
        return stats, summaryStrs

    stats, summaryStrs = _summarizeDets()
    return summaryStrs

def run_test(cfg, model, distributed, test_epoch=None):
    if distributed:
        model = model.module
    torch.cuda.empty_cache()  # TODO check if it helps
    iou_types = ("bbox",)
    if cfg.MODEL.MASK_ON:
        iou_types = iou_types + ("segm",)
    if cfg.MODEL.KEYPOINT_ON:
        iou_types = iou_types + ("keypoints",)
    output_folders = [None] * len(cfg.DATASETS.TEST)
    dataset_names = cfg.DATASETS.TEST
    if cfg.OUTPUT_DIR:
        for idx, dataset_name in enumerate(dataset_names):
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
            mkdir(output_folder)
            output_folders[idx] = output_folder
    data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
    for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
        inference_result = inference(
            model,
            data_loader_val,
            dataset_name=dataset_name,
            iou_types=iou_types,
            box_only=False if cfg.MODEL.FCOS_ON or cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY,
            device=cfg.MODEL.DEVICE,
            expected_results=cfg.TEST.EXPECTED_RESULTS,
            expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
            output_folder=output_folder,
        )
        synchronize()
        # import pdb; pdb.set_trace()
        summaryStrs = get_neat_inference_result(inference_result[2][0])
        # print('\n'.join(summaryStrs))
        summaryStrFinal = '\n'.join(summaryStrs)
        summaryStrFinal = '\n\nEpoch: ' + str(test_epoch) + '\n' + summaryStrFinal
        # with open(output_folder+'/summaryStrs.txt', 'w') as f_summaryStrs:
        with open(output_folder+'/summaryStrs.txt', 'a') as f_summaryStrs:
            f_summaryStrs.write(summaryStrFinal)

def main():
    parser = argparse.ArgumentParser(description="PyTorch Object Detection Training")
    parser.add_argument(
        "--config-file",
        default="maskfirst/mask_first.yaml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument(
        "--skip-test",
        dest="skip_test",
        help="Do not test the final model",
        action="store_true",
    )
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )

    args = parser.parse_args()

    num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    args.distributed = num_gpus > 1

    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(
            backend="nccl", init_method="env://"
        )
        synchronize()

    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    # cfg.merge_from_list(['OUTPUT_DIR', 'run/fcos_imprv_R_50_FPN_1x/vals_1'])
    cfg.freeze()

    output_dir = cfg.OUTPUT_DIR
    if output_dir:
        mkdir(output_dir)
        output_config_path = os.path.join(cfg.OUTPUT_DIR, 'new_config.yml')
        with open(output_config_path, 'w') as f:
            f.write(cfg.dump())

    logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank())
    logger.info("Using {} GPUs".format(num_gpus))
    logger.info(args)

    logger.info("Collecting env info (might take some time)")
    logger.info("\n" + collect_env_info())

    logger.info("Loaded configuration file {}".format(args.config_file))
    with open(args.config_file, "r") as cf:
        config_str = "\n" + cf.read()
        logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    model = train(cfg, args.local_rank, args.distributed)

    if not args.skip_test:
        run_test(cfg, model, args.distributed)

if __name__ == "__main__":
    main()

max_stars_repo_name: xianpf/CenterMask

max_issues_repo_name: xianpf/CenterMask

max_forks_repo_name: xianpf/CenterMask


text:
import cv2
import numpy as np
import os

class Segmentations():

    def __init__(self, dicom, model, params):
        self.oct_images = None
        self.dicom = dicom
        self.model = model
        self.params = params
        self.oct_segmentations = self.get_oct_and_segmentation()

    def save_segmentations(self, save_path):
        '''
        :param save_path: str; path were to save segmentations
        :return: None
        '''

        # create save dir if it does not exist
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # iterate through 49 segmentation's and save their plots separately
        for i in range( self.oct_segmentations.shape[0] ):

            # multiply by 255. to get clear visualization
            cv2.imwrite( os.path.join( save_path, str( i ) + ".png" ),
                         self.oct_segmentations[i, :, :] * 255. )

    def save_octs(self, save_path):
        '''
        :param save_path: str; path were to save octs
        :return: None
        '''

        # create save dir if it does not exist
        if not os.path.exists( save_path ):
            os.makedirs( save_path )

        # iterate through 49 oct's and save their plots separately
        for i in range( self.oct_images.shape[0] ):
            cv2.imwrite( os.path.join( save_path, str( i ) + ".png" ),
                         self.oct_images[i, :, :] )

    def get_oct_and_segmentation(self):
        '''
        :return: numpy array; segmented oct images
        '''

        # read in oct from dicom pixel array instance
        self.oct_images = self.dicom.dicom_file.pixel_array

        # predict segmentation masks for all oct images
        oct_segmentations = []
        for i in range( 0, self.oct_images.shape[0] ):

            # dimension of octs obtained with Heidelberg Spectrali's
            orig_height = 496
            orig_width = 512

            # model expect three channel input, stack oct if necessary
            if len(self.oct_images[i, :, :].shape) < 3:

                # stack oct image into 3 channel
                stacked_img = np.stack((self.oct_images[i, :, :],) * 3, axis = -1)

            else:
                stacked_img = self.oct_images[i, :, :]

            # resize and scale stacked image
            resized_image = cv2.resize( stacked_img, (self.params["img_shape"][0], self.params["img_shape"][0]),
                                        interpolation = cv2.INTER_NEAREST) / 255.

            # reshape image for prediction and predict
            reshaped_image = resized_image.reshape( 1, self.params["img_shape"][0], self.params["img_shape"][0], 3 )
            prediction = cv2.resize( self.model.predict(reshaped_image)[0, :, :, 0],
                                     (orig_width, orig_height),
                                     interpolation = cv2.INTER_NEAREST)

            # set class zero
            prediction[prediction < 0.5] = 0
            prediction[prediction >= 0.5] = 1

            oct_segmentations.append(prediction)

        oct_segmentations = np.array(oct_segmentations)
        return oct_segmentations

max_stars_repo_name: animesh/DeepRT

max_issues_repo_name: machikomomo/DeepRT

max_forks_repo_name: machikomomo/DeepRT


text:
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence
import numpy as np

from .lrcn_trainer import LRCNTrainer

class GVETrainer(LRCNTrainer):

    REQ_EVAL = True

    def __init__(self, args, model, dataset, data_loader, logger, device, checkpoint=None):
        super().__init__(args, model, dataset, data_loader, logger, device, checkpoint)
        self.rl_lambda = args.loss_lambda

    def train_step(self, image_input, word_inputs, word_targets, lengths,
            labels):
        # Forward, Backward and Optimize
        labels_onehot = self.model.convert_onehot(labels)
        labels_onehot = labels_onehot.to(self.device)
        self.model.zero_grad()
        outputs = self.model(image_input, word_inputs, lengths, labels,
                labels_onehot=labels_onehot)

        # Reinforce loss
        # Sample sentences
        sample_ids, log_ps, lengths = self.model.generate_sentence(image_input, self.start_word,
                self.end_word, labels, labels_onehot=labels_onehot, max_sampling_length=50, sample=True)
        # Order sampled sentences/log_probabilities/labels by sentence length (required by LSTM)
        lengths = lengths.cpu().numpy()
        sort_idx = np.argsort(-lengths)
        lengths = lengths[sort_idx]
        sort_idx = torch.tensor(sort_idx, device=self.device, dtype=torch.long)
        labels = labels[sort_idx]
        labels = labels.to(self.device)
        log_ps = log_ps[sort_idx,:]
        sample_ids = sample_ids[sort_idx,:]

        class_pred = self.model.sentence_classifier(sample_ids, lengths)
        class_pred = F.softmax(class_pred, dim=1)
        rewards = class_pred.gather(1, labels.view(-1,1)).squeeze()
        r_loss = -(log_ps.sum(dim=1) * rewards).sum()

        loss = self.rl_lambda * r_loss/labels.size(0) + self.criterion(outputs, word_targets)
        loss.backward()
        #nn.utils.clip_grad_norm(self.params, 10)
        self.optimizer.step()

        return loss


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) 2020, NVIDIA CORPORATION. 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.

"""Losses used for Mask-RCNN."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from distutils.version import LooseVersion

import tensorflow as tf
from sagemakercv.core import box_utils
from .losses import _sigmoid_cross_entropy, _huber_loss, _giou_loss, _ciou_loss, _l1_loss

DEBUG_LOSS_IMPLEMENTATION = False

if LooseVersion(tf.__version__) < LooseVersion("2.0.0"):
    from tensorflow.python.keras.utils import losses_utils
    ReductionV2 = losses_utils.ReductionV2
else:
    ReductionV2 = tf.keras.losses.Reduction

class RPNLoss(object):

    def __init__(self, min_level=2,
                       max_level=6,
                       box_loss_type='huber',
                       train_batch_size_per_gpu=1,
                       rpn_batch_size_per_im=256,
                       label_smoothing=0.0,
                       rpn_box_loss_weight=1.0):
        self.min_level = min_level
        self.max_level = max_level
        self.box_loss_type = box_loss_type
        self.train_batch_size_per_gpu = train_batch_size_per_gpu
        self.rpn_batch_size_per_im = rpn_batch_size_per_im
        self.label_smoothing = label_smoothing
        self.rpn_box_loss_weight = rpn_box_loss_weight

    def _rpn_box_loss(self, box_outputs, box_targets, normalizer=1.0, delta=1. / 9):
        """Computes box regression loss."""
        # delta is typically around the mean value of regression target.
        # for instances, the regression targets of 512x512 input with 6 anchors on
        # P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].

        with tf.name_scope('rpn_box_loss'):
            mask = tf.not_equal(box_targets, 0.0)
            mask = tf.cast(mask, tf.float32)

            assert mask.dtype == tf.float32

            # The loss is normalized by the sum of non-zero weights before additional
            # normalizer provided by the function caller.
            if self.box_loss_type == 'huber':
                box_loss = _huber_loss(y_true=box_targets, y_pred=box_outputs, weights=mask, delta=delta)
                # TODO: Test giou loss for rpn head
                '''elif self.box_loss_type == 'giou':
                    box_loss = _giou_loss(y_true=box_targets, y_pred=box_outputs, weights=mask)
                elif self.box_loss_type == 'ciou':
                    box_loss = _ciou_loss(y_true=box_targets, y_pred=box_outputs, weights=mask)'''
            elif self.box_loss_type == 'l1_loss':
                box_loss = _l1_loss(y_true=box_targets, y_pred=box_outputs, weights=mask, delta=delta)
            else:
                # box_loss = _l1_loss(y_true=box_targets, y_pred=box_outputs, weights=mask, delta=delta)
                raise NotImplementedError
            # box_loss = _l1_loss(y_true=box_targets, y_pred=box_outputs, weights=mask, delta=delta)

            assert box_loss.dtype == tf.float32

            if isinstance(normalizer, tf.Tensor) or normalizer != 1.0:
                box_loss /= normalizer

            assert box_loss.dtype == tf.float32

        return box_loss

    def _rpn_score_loss(self, score_outputs, score_targets, normalizer=1.0):
        """Computes score loss."""

        with tf.name_scope('rpn_score_loss'):

            # score_targets has three values:
            # * (1) score_targets[i]=1, the anchor is a positive sample.
            # * (2) score_targets[i]=0, negative.
            # * (3) score_targets[i]=-1, the anchor is don't care (ignore).

            mask = tf.math.greater_equal(score_targets, 0)
            mask = tf.cast(mask, dtype=tf.float32)

            score_targets = tf.maximum(score_targets, tf.zeros_like(score_targets))
            score_targets = tf.cast(score_targets, dtype=tf.float32)

            assert score_outputs.dtype == tf.float32
            assert score_targets.dtype == tf.float32

            score_loss = _sigmoid_cross_entropy(
                multi_class_labels=score_targets,
                logits=score_outputs,
                weights=mask,
                sum_by_non_zeros_weights=False,
                label_smoothing=self.label_smoothing
            )

            assert score_loss.dtype == tf.float32

            if isinstance(normalizer, tf.Tensor) or normalizer != 1.0:
                score_loss /= normalizer

            assert score_loss.dtype == tf.float32

        return score_loss

    def __call__(self, score_outputs, box_outputs, labels):
        with tf.name_scope('rpn_loss'):
            score_losses = []
            box_losses = []
            for level in range(int(self.min_level), int(self.max_level + 1)):
                score_targets_at_level = labels['score_targets_%d' % level]
                box_targets_at_level = labels['box_targets_%d' % level]
                score_losses.append(
                    self._rpn_score_loss(
                        score_outputs=score_outputs[level],
                        score_targets=score_targets_at_level,
                        normalizer=tf.cast(self.train_batch_size_per_gpu \
                                           * self.rpn_batch_size_per_im, dtype=tf.float32)
                    ))
                box_losses.append(self._rpn_box_loss(
                    box_outputs=box_outputs[level],
                    box_targets=box_targets_at_level,
                    normalizer=1.0
                ))
            # Sum per level losses to total loss.
            rpn_score_loss = tf.add_n(score_losses)
            rpn_box_loss = self.rpn_box_loss_weight * tf.add_n(box_losses)
            total_rpn_loss = rpn_score_loss + rpn_box_loss
        return total_rpn_loss, rpn_score_loss, rpn_box_loss

max_stars_repo_name: aws-samples/amazon-sagemaker-cv

max_issues_repo_name: aws-samples/amazon-sagemaker-cv

max_forks_repo_name: aws-samples/amazon-sagemaker-cv

reason: from sage