text:
import torch
from torch import nn, optim
import json
from .data_loader import SentenceRELoader, BagRELoader
from .utils import AverageMeter
from tqdm import tqdm
import sys
sys.path.append("..")
import HiCLRE.hiclre.framework.data_loader as data_load_py
from transformers import BertModel, BertTokenizer
from scipy.spatial.distance import cdist
import torch.nn.functional as F

pretrain_path='/pythonCode/HiCLRE/pretrain/bert-base-uncased'
class BagRE(nn.Module):

    def __init__(self,
                 model,
                 train_path,
                 val_path,
                 test_path,
                 ckpt,
                 batch_size=32,
                 max_epoch=100,
                 lr=0.1,
                 weight_decay=1e-5,
                 opt='sgd',
                 bag_size=0,
                 tau=0.5,
                 lambda_1=0.4,
                 lambda_2=0.4,
                 lambda_3=0.1,
                 lambda_4=0.1,
                 loss_weight=False):

        super().__init__()
        self.max_epoch = max_epoch
        self.bag_size = bag_size
        self.batch_size = batch_size
        self.tokenizer = BertTokenizer.from_pretrained(pretrain_path)
        self.fc = nn.Linear(model.sentence_encoder.hidden_size, model.num_class)
        self.diag = nn.Parameter(torch.ones(model.sentence_encoder.hidden_size))
        self.softmax = nn.Softmax(-1)
        self.drop = nn.Dropout()
        self.fc = nn.Linear(model.sentence_encoder.hidden_size, len(model.rel2id))
        self.max_length=128
        self.tau = tau
        self.lambda_1 = lambda_1
        self.lambda_2 = lambda_2
        self.lambda_3 = lambda_3
        self.lambda_4 = lambda_4
        # Load data
        if train_path != None:
            self.train_loader = BagRELoader(
                train_path,
                model.rel2id,
                model.sentence_encoder.tokenize,
                batch_size,
                True,
                # False,
                bag_size=bag_size,
                entpair_as_bag=False)

        if val_path != None:
            self.val_loader = BagRELoader(
                val_path,
                model.rel2id,
                model.sentence_encoder.tokenize,
                batch_size,
                False,
                bag_size=bag_size,
                entpair_as_bag=True)

        if test_path != None:
            self.test_loader = BagRELoader(
                test_path,
                model.rel2id,
                model.sentence_encoder.tokenize,
                batch_size,
                False,
                bag_size=bag_size,
                entpair_as_bag=True
            )
        # Model
        self.model = nn.DataParallel(model)
        # Criterion
        if loss_weight:
            self.criterion = nn.CrossEntropyLoss(weight=self.train_loader.dataset.weight)
        else:
            self.criterion = nn.CrossEntropyLoss()
        # Params and optimizer
        params = self.model.parameters()
        self.lr = lr
        if opt == 'sgd':
            self.optimizer = optim.SGD(params, lr, weight_decay=weight_decay)
        elif opt == 'adam':
            self.optimizer = optim.Adam(params, lr, weight_decay=weight_decay)
        elif opt == 'adamw':
            from transformers import AdamW
            params = list(self.named_parameters())
            no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
            grouped_params = [
                {
                    'params': [p for n, p in params if not any(nd in n for nd in no_decay)],
                    'weight_decay': 0.01,
                    'lr': lr,
                    'ori_lr': lr
                },
                {
                    'params': [p for n, p in params if any(nd in n for nd in no_decay)],
                    'weight_decay': 0.0,
                    'lr': lr,
                    'ori_lr': lr
                }
            ]
            self.optimizer = AdamW(grouped_params, correct_bias=False)
        else:
            raise Exception("Invalid optimizer. Must be 'sgd' or 'adam' or 'bert_adam'.")
        # Cuda
        if torch.cuda.is_available():
            self.cuda()
        # Ckpt
        self.ckpt = ckpt

    def train_model(self, metric='auc'):
        best_metric = 0
        sentence_rep_memory = {}
        entity_rep_memory = {}
        align_list = []
        uniform_list = []

        for epoch in range(self.max_epoch):
            # Train
            self.train()
            print("=== Epoch %d train ===" % epoch)
            avg_loss = AverageMeter()
            avg_acc = AverageMeter()
            avg_pos_acc = AverageMeter()

            # inner_bag_rep_memory = {}
            inner_sentence_rep_memory = {}
            inner_entity_rep_memory = {}
            num_no_key = 0

            t = tqdm(self.train_loader)
            t_len= t.__len__()
            for iter, data in enumerate(t):
                if torch.cuda.is_available():
                    for i in range(len(data)):
                        try:
                            data[i] = data[i].cuda()
                        except:
                            pass
                label = data[0]
                bag_name = data[1]
                scope = data[2]
                args = data[3:7]
                head_end = data[7]
                tail_end = data[8]
                sentence_index = data[9]
                sentence_index_list = (torch.squeeze(sentence_index.view(self.batch_size * self.bag_size,-1))).detach().cpu().numpy()

                logits, rep, rep_head, rep_tail, bag_perturb, sent_perturb, entity_perturb = self.model( head_end, tail_end, label, scope, *args, bag_size=self.bag_size)

                #---------------------store sentence rep to memory in sentIndex order---------------------------
                rep_list = ((rep.view(self.batch_size * self.bag_size,-1)).detach().cpu().numpy()).tolist()
                sentenceIndex_rep_dic = dict(map(lambda x, y: [x, y], sentence_index_list, rep_list))
                inner_sentence_rep_memory.update(sentenceIndex_rep_dic)
                # ---------------------store sentence rep to memory in sentIndex order---------------------------

                # ---------------------store entity rep to memory in sentIndex order---------------------------
                rep_entity = torch.cat([rep_head, rep_tail], 1)
                rep_entity_list = (rep_entity.detach().cpu().numpy()).tolist()
                sentenceIndex_rep_entity_dic = dict(map(lambda x, y: [x, y], sentence_index_list, rep_entity_list))
                inner_entity_rep_memory.update(sentenceIndex_rep_entity_dic)
                # ---------------------store entity rep to memory in sentIndex order---------------------------

                loss_crossEnropy = self.criterion(logits, label)
                loss_crossEnropy.detach()

                # --------------------sentence-level ---------------------

                rep = rep.view(self.batch_size * self.bag_size, -1)
                sent_perturb = sent_perturb.view(self.batch_size * self.bag_size, -1)
                if epoch != 0:
                    last_epoch_rep_list = []
                    for i in range(len(sentence_index_list)):

                        if sentence_index_list[i] in sentence_rep_memory.keys() :
                            last_epoch_rep_list.append(sentence_rep_memory[sentence_index_list[i]])
                        else:
                            num_no_key += 1
                            last_epoch_rep_list.append(torch.zeros(1536).detach().cpu().numpy().tolist())
                    last_epoch_rep_tensor = torch.tensor(last_epoch_rep_list)

                    sentence_simi = torch.cosine_similarity( rep, last_epoch_rep_tensor.to('cuda'))
                    sentence_simi = torch.unsqueeze(sentence_simi, dim=1)
                    sent_sim_expand = ((self.max_epoch - epoch) / self.max_epoch) * sentence_simi / torch.norm(sentence_simi, 2)
                    sent_perturb = sent_perturb + sent_sim_expand

                rep_after_perturb = rep + sent_perturb
                sim_cos = torch.cosine_similarity(rep, rep_after_perturb)

                sim_self = torch.cosine_similarity(rep.unsqueeze(1), rep.unsqueeze(0), dim=2)
                loss_sentence = - torch.log((torch.exp(sim_cos / self.tau)) / torch.exp(sim_self / self.tau).sum(0))
                # --------------------sentence-level ---------------------

                #--------------------bag-level ---------------------
                if epoch != 0:
                    batch_size = label.size(0)
                    query = label.unsqueeze(1)
                    att_mat = self.fc.weight[query]
                    att_mat = att_mat * self.diag.unsqueeze(0)
                    last_epoch_rep_tensor = last_epoch_rep_tensor.view(self.batch_size, self.bag_size, -1).to('cuda')
                    att_score = (last_epoch_rep_tensor * att_mat).sum(-1)
                    softmax_att_score = self.softmax(att_score)
                    last_bag_rep_3dim = (softmax_att_score.unsqueeze(-1) * last_epoch_rep_tensor)
                    last_bag_rep = last_bag_rep_3dim.sum(1)
                    last_bag_rep = self.drop(last_bag_rep)
                    last_bag_logits = self.fc(last_bag_rep)

                    bag_simi = torch.cosine_similarity( logits, last_bag_logits )
                    bag_simi = torch.unsqueeze(bag_simi, dim=1)

                    bag_sim_expand = ((self.max_epoch - epoch) / self.max_epoch ) * bag_simi / torch.norm(bag_simi, 2)
                    bag_perturb = bag_perturb + bag_sim_expand

                logits_after_perturb = logits + bag_perturb
                simcse_adv_cos = torch.cosine_similarity(logits, logits_after_perturb)
                adv_sim_self = torch.cosine_similarity(logits.unsqueeze(1), logits.unsqueeze(0), dim=2)
                loss_bag = - torch.log((torch.exp(simcse_adv_cos / self.tau)) / torch.exp(adv_sim_self / self.tau).sum(0))

                #----------------align and uniform-------------------------
                align = (sim_cos) / (self.batch_size * self.bag_size)
                align = align.detach().cpu().numpy().tolist()
                align_list = align_list + align
                uniform = (-  2 * sim_self).sum(0) / (self.batch_size * self.bag_size)
                uniform = uniform.detach().cpu().numpy().tolist()
                uniform_list = uniform_list + uniform
                # ----------------align and uniform------------------------

                # --------------------bag-level ---------------------

                # --------------------entity-level ---------------------
                rep_entity = rep_entity.view(self.batch_size * self.bag_size, -1)
                entity_perturb = entity_perturb.view(self.batch_size * self.bag_size, -1)
                if epoch != 0:
                    last_epoch_rep_entity_list = []
                    for i in range(len(sentence_index_list)):

                        if sentence_index_list[i] in entity_rep_memory.keys():
                            last_epoch_rep_entity_list.append(entity_rep_memory[sentence_index_list[i]])
                        else:
                            last_epoch_rep_entity_list.append(torch.zeros(1536).detach().cpu().numpy().tolist())
                    last_epoch_rep_entity_tensor = torch.tensor(last_epoch_rep_entity_list)

                    entity_simi = torch.cosine_similarity(rep_entity, last_epoch_rep_entity_tensor.to('cuda'))
                    entity_simi = torch.unsqueeze(entity_simi, dim=1)
                    entity_sim_expand = ((self.max_epoch - epoch) / self.max_epoch) * entity_simi / torch.norm(entity_simi, 2)
                    entity_perturb = entity_perturb + entity_sim_expand
                entity_rep_after_perturb = rep_entity + entity_perturb
                entity_sim_cos = torch.cosine_similarity(rep_entity, entity_rep_after_perturb)
                entity_sim_self = torch.cosine_similarity(rep_entity.unsqueeze(1), rep_entity.unsqueeze(0), dim=2)
                loss_entity = - torch.log((torch.exp(entity_sim_cos / self.tau)) / torch.exp(entity_sim_self / self.tau).sum(0))
                # --------------------entity-level ---------------------

                mean_loss_bag = (torch.sum(loss_bag, dim=0)) / (self.batch_size * self.bag_size)
                mean_loss_sentence = (torch.sum(loss_sentence, dim=0)) / (self.batch_size * self.bag_size) / 4
                mean_loss_entity = (torch.sum(loss_entity, dim=0)) / (self.batch_size * self.bag_size) / 4

                total_loss = self.lambda_1 *loss_crossEnropy + self.lambda_2*mean_loss_bag + self.lambda_3*mean_loss_sentence + self.lambda_4*mean_loss_entity

                score, pred = logits.max(-1)
                acc = float((pred == label).long().sum()) / label.size(0)

                pos_total = (label != 0).long().sum()
                pos_correct = ((pred == label).long() * (label != 0).long()).sum()
                if pos_total > 0:
                    pos_acc = float(pos_correct) / float(pos_total)
                else:
                    pos_acc = 0

                avg_loss.update(total_loss.item(), 1)
                avg_acc.update(acc, 1)
                avg_pos_acc.update(pos_acc, 1)
                t.set_postfix(loss=avg_loss.avg, acc=avg_acc.avg, pos_acc=avg_pos_acc.avg)

                total_loss.backward()
                self.optimizer.step()
                self.optimizer.zero_grad()

                #--------------------add more test operation in each epoch------------------------
                if iter % (int((t_len / 3) + 1)) == 0 and iter > 0:
                    val_index = (iter // (int((t_len / 3) + 1))) - 1
                    print("=== Epoch %d-%d val === " % (epoch,  val_index))
                    result = self.eval_model(self.val_loader)
                    print("AUC: %.4f" % result['auc'])
                    print("Micro F1: %.4f" % (result['max_micro_f1']))
                    if result[metric] > best_metric:
                        print("Best ckpt and saved.")
                        torch.save({'state_dict': self.model.module.state_dict()}, self.ckpt)
                        best_metric = result[metric]
                #--------------------add more test operation in each epoch------------------------
            sentence_rep_memory.update( inner_sentence_rep_memory )
            entity_rep_memory.update( inner_entity_rep_memory )

            # Val
            print("=== Epoch %d-2 val ===" % epoch)
            result = self.eval_model(self.val_loader)
            print("AUC: %.4f" % result['auc'])
            print("Micro F1: %.4f" % (result['max_micro_f1']))
            if result[metric] > best_metric:
                print("Best ckpt and saved.")
                torch.save({'state_dict': self.model.module.state_dict()}, self.ckpt)
                best_metric = result[metric]

        print("Best %s on val set: %f" % (metric, best_metric))
        # return align_list, uniform_list

    def eval_model(self, eval_loader):
        self.model.eval()
        with torch.no_grad():
            t = tqdm(eval_loader)
            pred_result = []
            for iter, data in enumerate(t):
                if torch.cuda.is_available():
                    for i in range(len(data)):
                        try:
                            data[i] = data[i].cuda()
                        except:
                            pass
                label = data[0]
                bag_name = data[1]
                scope = data[2]
                args = data[3:7]
                head_end = data[7]
                tail_end = data[8]
                sentence_index = data[9]
                logits, rep, rep_head, rep_tail, bag_perturb, sent_perturb, entity_perturb = self.model(head_end, tail_end,  None, scope, *args, train=False, bag_size=self.bag_size) # results after softmax

                logits = logits.cpu().numpy()
                for i in range(len(logits)):
                    for relid in range(self.model.module.num_class):
                        if self.model.module.id2rel[relid] != 'NA':
                            pred_result.append({
                                'entpair': bag_name[i][:2],
                                'relation': self.model.module.id2rel[relid],
                                'score': logits[i][relid],
                                'sentence_index': sentence_index[i]
                            })
            result = eval_loader.dataset.eval(pred_result)
        return result

    def load_state_dict(self, state_dict):
        self.model.module.load_state_dict(state_dict)

max_stars_repo_name: MatNLP/HiCLRE

max_issues_repo_name: MatNLP/HiCLRE

max_forks_repo_name: MatNLP/HiCLRE

reason: from scipy


text:
#! /usr/bin/env python

import numpy as np
import pandas as pd
from scipy.interpolate import RegularGridInterpolator

def get_df_interp_func(filename=None, df=None, gauss=True, mm=False, scipy_interp=False, bounds=None):
    '''
    This factory function will return an interpolating function for any field map. An input x,y,z will output the corresponding Bx,By,Bz or Br,Bphi,Bz. Will decide later if linear interpolation is good enough.

    Assumed file for input has length: meters, Bfield: Gauss
    '''
    # load dataframe if not passed in
    if df is None:
        if ".p" in filename:
            df = pd.read_pickle(filename)
        else:
            df = pd.read_csv(filename)
    else:
        df = df.copy()
    if not gauss:
        df["Bx"] = df["Bx"] / 1e4
        df["By"] = df["By"] / 1e4
        df["Bz"] = df["Bz"] / 1e4
    if mm:
        df["X"] = df["X"] * 1e3
        df["Y"] = df["Y"] * 1e3
        df["Z"] = df["Z"] * 1e3

    xs = df.X.unique()
    ys = df.Y.unique()
    zs = df.Z.unique()

    if bounds is not None:
        xmin = xs[xs < bounds.xmin][-1]
        xmax = xs[xs > bounds.xmax][0]
        ymin = ys[ys < bounds.ymin][-1]
        ymax = ys[ys > bounds.ymax][0]
        zmin = zs[zs < bounds.zmin][-1]
        zmax = zs[zs > bounds.zmax][0]
        query_string = f"X>={xmin} & X<={xmax} & Y>={ymin} & Y<={ymax} & Z>={zmin} & Z<={zmax}"
        df = df.query(query_string)

        xs = df.X.unique()
        ys = df.Y.unique()
        zs = df.Z.unique()

    dx = xs[1]-xs[0]
    dy = ys[1]-ys[0]
    dz = zs[1]-zs[0]

    lx = len(xs)
    ly = len(ys)
    lz = len(zs)

    df_np = df[["X","Y","Z","Bx","By","Bz"]].values

    x, y, z, bx, by, bz = df_np.T
    BX = bx.reshape(lx, ly, lz)
    BY = by.reshape(lx, ly, lz)
    BZ = bz.reshape(lx, ly, lz)

    interp_func_Bx = RegularGridInterpolator((xs, ys, zs), BX)
    interp_func_By = RegularGridInterpolator((xs, ys, zs), BY)
    interp_func_Bz = RegularGridInterpolator((xs, ys, zs), BZ)

    def scipy_interp_func(pts):
        Bx = interp_func_Bx(pts.T)
        By = interp_func_By(pts.T)
        Bz = interp_func_Bz(pts.T)
        return np.array([Bx, By, Bz])

    def get_cube(x, y, z):
        a_x, a_y, a_z = len(xs[xs <= x]) - 1, len(ys[ys <= y]) - 1, len(zs[zs <= z]) - 1
        corner_a = (ly * lz) * a_x + (lz) * a_y + a_z
        corner_b = corner_a + lz
        corner_c = corner_a + ly * lz
        corner_d = corner_a + ly * lz + lz
        index_list = [corner_a,corner_a+1,corner_b,corner_b+1,
        corner_c,corner_c+1,corner_d,corner_d+1]
        return df_np[index_list]

    def interp_single(xd,yd,zd,ff):
        c00 = ff[0,0,0]*(1 - xd) + ff[1,0,0] * xd
        c01 = ff[0,0,1]*(1 - xd) + ff[1,0,1] * xd
        c10 = ff[0,1,0]*(1 - xd) + ff[1,1,0] * xd
        c11 = ff[0,1,1]*(1 - xd) + ff[1,1,1] * xd

        c0 = c00 * (1 - yd) + c10 * yd
        c1 = c01 * (1 - yd) + c11 * yd
        return c0 * (1 - zd) + c1 * zd

    def interp(p_vec):
        cube = get_cube(*p_vec)

        xx = np.unique(cube[:,0])
        yy = np.unique(cube[:,1])
        zz = np.unique(cube[:,2])

        bxs_grid = cube[:,3].reshape((2,2,2))
        bys_grid = cube[:,4].reshape((2,2,2))
        bzs_grid = cube[:,5].reshape((2,2,2))

        xd = (p_vec[0]-xx[0])/(xx[1]-xx[0])
        yd = (p_vec[1]-yy[0])/(yy[1]-yy[0])
        zd = (p_vec[2]-zz[0])/(zz[1]-zz[0])

        bx = interp_single(xd,yd,zd, bxs_grid)
        by = interp_single(xd,yd,zd, bys_grid)
        bz = interp_single(xd,yd,zd, bzs_grid)

        return np.array([bx,by,bz])

    if scipy_interp:
        return scipy_interp_func
    else:
        return interp

ext: py

reason: import numpy,from scipy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May  5 10:42:50 2022

@author: lukepinkel
"""

import numpy as np
from .special_mats import lmat, nmat
from .numerical_derivs import jac_approx
from .dchol import dchol, unit_matrices

def _hess_chain_rule(d2z_dy2, dy_dx, dz_dy, d2y_dx2):
    H1 = np.einsum("ijk,kl->ijl", d2z_dy2, dy_dx, optimize=True)
    H1 = np.einsum("ijk,jl->ilk", H1, dy_dx, optimize=True)
    H2 = np.einsum("ij,jkl->ikl", dz_dy, d2y_dx2, optimize=True)
    d2z_dx2 = H1 + H2
    return d2z_dx2

def _vecl(x):
    m = x.shape[-1]
    ix, jx = np.triu_indices(m, k=1)
    res = x[...,jx, ix]
    return res

def _invecl(x):
    old_shape = x.shape
    n = x.shape[-1]
    m = int((np.sqrt(8 * n + 1) + 1) // 2)
    out_shape = old_shape[:-1] + (m, m)
    res = np.zeros(out_shape, dtype=x.dtype)
    ix, jx = np.triu_indices(m, k=1)
    is_diag = ix == jx
    off_diag_elem = x[...,~is_diag]
    ixo, jxo = ix[~is_diag], jx[~is_diag]
    ixd = jxd = np.arange(m)
    res[..., jxo, ixo] = off_diag_elem
    res[..., ixo, jxo] = off_diag_elem
    res[..., ixd, jxd] = 1
    return res

def _vech(x):
    m = x.shape[-1]
    ix, jx = np.triu_indices(m, k=0)
    res = x[...,jx, ix]
    return res

def _invech(x):
    old_shape = x.shape
    n = x.shape[-1]
    m = int((np.sqrt(8 * n + 1) - 1) // 2)
    out_shape = old_shape[:-1] + (m, m)
    res = np.zeros(out_shape, dtype=x.dtype)
    ix, jx = np.triu_indices(m, k=0)
    is_diag = ix == jx
    diag_elements = x[..., is_diag]
    off_diag_elem = x[...,~is_diag]
    ixo, jxo = ix[~is_diag], jx[~is_diag]
    ixd, jxd = ix[is_diag],  jx[is_diag]
    res[..., jxo, ixo] = off_diag_elem
    res[..., ixo, jxo] = off_diag_elem
    res[..., ixd, jxd] = diag_elements
    return res

def lhv_size_to_mat_size(lhv_size):
    mat_size = int((np.sqrt(8 * lhv_size + 1) + 1) // 2)
    return mat_size

def mat_size_to_lhv_size(mat_size):
    lhv_size = int(mat_size * (mat_size - 1) // 2)
    return lhv_size

def hv_size_to_mat_size(lhv_size):
    mat_size = int((np.sqrt(8 * lhv_size + 1) - 1) // 2)
    return mat_size

def mat_size_to_hv_size(mat_size):
    lhv_size = int(mat_size * (mat_size + 1) // 2)
    return lhv_size

def lower_half_vec(x):
    mat_size = x.shape[-1]
    i, j = np.triu_indices(mat_size, k=1)
    y = x[...,j, i]
    return y

def inv_lower_half_vec(y):
    old_shape = y.shape
    lhv_size = y.shape[-1]
    mat_size = lhv_size_to_mat_size(lhv_size)
    out_shape = old_shape[:-1] + (mat_size, mat_size)
    x = np.zeros(out_shape, dtype=y.dtype)
    i, j = np.triu_indices(mat_size, k=1)
    x[..., j, i] = y
    return x

def lhv_indices(shape):
    arr_inds = np.indices(shape)
    lhv_inds = [lower_half_vec(x) for x in arr_inds]
    return lhv_inds

def hv_indices(shape):
    arr_inds = np.indices(shape)
    hv_inds = [_vech(x) for x in arr_inds]
    return hv_inds

def lhv_ind_parts(mat_size):
    i = np.cumsum(np.arange(mat_size))
    return list(zip(*(i[:-1], i[1:])))

def lhv_row_norms(y):
    lhv_size = y.shape[-1]
    mat_size = lhv_size_to_mat_size(lhv_size)
    r, c = lhv_indices((mat_size, mat_size))
    rj = np.argsort(r)
    ind_partitions = lhv_ind_parts(mat_size)
    row_norms = np.zeros(mat_size)
    row_norms[0] = 1.0
    for i, (a, b) in enumerate(ind_partitions):
        ii = rj[a:b]
        row_norms[i+1] = np.sqrt(np.sum(y[ii]**2)+1)
    return row_norms

class UnconstrainedCholeskyCorr(object):

    def __init__(self, mat_size):
        self.mat_size = mat_size
        self.lhv_size = mat_size_to_lhv_size(mat_size)
        self.row_inds, self.col_inds = lhv_indices((mat_size, mat_size))
        self.row_sort = np.argsort(self.row_inds)
        self.ind_parts = lhv_ind_parts(mat_size)
        self.row_norm_inds = [self.row_sort[a:b] for a, b in self.ind_parts]

    def _fwd(self, x):
        row_norms = np.zeros_like(x)
        for ii in self.row_norm_inds:
            row_norms[ii] = np.sqrt(np.sum(x[ii]**2)+1)
        y = x / row_norms
        return y

    def _rvs(self, y):
        diag = np.zeros_like(y)
        for ii in self.row_norm_inds:
            diag[ii] = np.sqrt(1-np.sum(y[ii]**2))
        x = y / diag
        return x

    def _jac_fwd(self, x):
        dy_dx = np.zeros((x.shape[0],)*2)
        for ii in self.row_norm_inds:
            xii = x[ii]
            s = np.sqrt(np.sum(xii**2)+1)
            v1 = 1.0 / s * np.eye(len(ii))
            v2 = 1.0 / (s**3) * xii[:, None] * xii[:,None].T
            dy_dx[ii, ii[:, None]] = v1 - v2
        return dy_dx

    def _hess_fwd(self, x):
        d2y_dx2 = np.zeros((x.shape[0],)*3)
        for ii in self.row_norm_inds:
            x_ii = x[ii]
            s = np.sqrt(1.0 + np.sum(x_ii**2))
            s3 = s**3
            s5 = s**5
            for i in ii:
                for j in ii:
                    for k in ii:
                        t1 = -1.0*(j==k) / s3 * x[i]
                        t2 = -1.0*(j==i) / s3 * x[k]
                        t3 = -1.0*(k==i) / s3 * x[j]
                        t4 = 3.0 / (s5) * x[j] * x[k] * x[i]
                        d2y_dx2[i, j, k] = t1+t2+t3+t4
        return d2y_dx2

    def _jac_rvs(self, y):
        dx_dy = np.zeros((y.shape[0],)*2)
        for ii in self.row_norm_inds:
            yii = y[ii]
            s = np.sqrt(1.0 - np.sum(yii**2))
            v1 = 1.0 / s * np.eye(len(ii))
            v2 = 1.0 / (s**3) * yii[:, None] * yii[:,None].T
            dx_dy[ii, ii[:, None]] = v1 + v2
        return dx_dy

    def _hess_rvs(self, y):
        d2x_dy2 = np.zeros((y.shape[0],)*3)
        for ii in self.row_norm_inds:
            y_ii = y[ii]
            s = np.sqrt(1.0 - np.sum(y_ii**2))
            s3 = s**3
            s5 = s**5
            for i in ii:
                for j in ii:
                    for k in ii:
                        t1 = 1.0*(j==k) / s3 * y[i]
                        t2 = 1.0*(j==i) / s3 * y[k]
                        t3 = 1.0*(k==i) / s3 * y[j]
                        t4 = 3.0 / s5 * y[i] * y[j] * y[k]
                        d2x_dy2[i, j, k] = t1 + t2 + t3 + t4
        return d2x_dy2

class CovCorr(object):

    def __init__(self, mat_size):
        self.mat_size = mat_size
        self.hv_size = mat_size_to_hv_size(mat_size)
        self.row_inds, self.col_inds = hv_indices((mat_size, mat_size))
        self.diag_inds, = np.where(self.row_inds==self.col_inds)
        self.tril_inds, = np.where(self.row_inds!=self.col_inds)
        self.row_diag_inds = self.diag_inds[self.row_inds]
        self.col_diag_inds = self.diag_inds[self.col_inds]
        ii = self.row_diag_inds!=self.col_diag_inds
        self.dr_perm = np.vstack((self.col_diag_inds[ii], self.row_diag_inds[ii])).T.flatten()
        self.dc_perm = np.repeat(self.tril_inds, 2)
        self.ii = ii

    def _fwd(self, x):
        sj = np.sqrt(x[self.row_diag_inds])
        sk = np.sqrt(x[self.col_diag_inds])
        y = x / (sj * sk)
        y[self.diag_inds] = x[self.diag_inds] #np.log(x[self.diag_inds])
        return y

    def _rvs(self, y):
        sj = np.sqrt(y[self.row_diag_inds])
        sk = np.sqrt(y[self.col_diag_inds])
        x = y * sj * sk
        x[self.diag_inds] = y[self.diag_inds]
        return x

    def _jac_fwd(self, x):
        dy_dx = np.zeros((x.shape[0],)*2)
        sj = np.sqrt(x[self.row_diag_inds])
        sk = np.sqrt(x[self.col_diag_inds])
        t1 = 1 / (sj*sk)
        t2 = -1.0 / 2.0 * x / (sj**3 * sk)
        t3 = -1.0 / 2.0 * x / (sk**3 * sj)
        t = np.vstack((t3, t2)).T[self.ii].flatten()
        dy_dx[self.diag_inds, self.diag_inds] = 1.0# / x[self.diag_inds]
        dy_dx[self.tril_inds, self.tril_inds] = t1[self.tril_inds]
        dy_dx[self.dc_perm, self.dr_perm] = t
        return dy_dx

    def _jac_rvs(self, y):
        dx_dy = np.zeros((y.shape[0],)*2)
        sj = np.sqrt(y[self.row_diag_inds])
        sk = np.sqrt(y[self.col_diag_inds])
        t1 = (sj*sk)
        t2 = 1.0 / 2.0 * y * sk / sj
        t3 = 1.0 / 2.0 * y * sj / sk
        t = np.vstack((t3, t2)).T[self.ii].flatten()
        dx_dy[self.diag_inds, self.diag_inds] = 1.0#np.exp(y[self.diag_inds])
        dx_dy[self.tril_inds, self.tril_inds] = t1[self.tril_inds]
        dx_dy[self.dc_perm, self.dr_perm] = t
        return dx_dy

    def _hess_fwd(self, x):
        rix, cix = self.row_diag_inds, self.col_diag_inds
        d2y_dx2 = np.zeros((self.hv_size,)*3)
        for i in range(self.hv_size):
            j, k = rix[i], cix[i]
            if j!=k:
                xi, xj, xk = x[i], x[j], x[k] #w, y, z = x[i], x[j], x[k]
                sjk =  np.sqrt(xj * xk)       #syz = np.sqrt(y * z)
                d2yi_dxidxj = -xk / (2.0 * sjk**3) #d2y_dwdy = -z / (2.0 * syz**3)
                d2yi_dxidxk = -xj / (2.0 * sjk**3) #d2y_dwdz = -y / (2.0 * syz**3)
                d2yi_dxjdxj = (3.0 * xi * xk**2)   / (4.0 * sjk**5)  #d2y_dydy = (3.0 * w * z**2)  / (4.0 * syz**5)
                d2yi_dxjdxk = (3.0 * xi * xj * xk) / (4.0 * sjk**5) - xi / (2.0 * sjk**3) #d2y_dydz = (3.0 * w * y * z) / (4.0 * syz**5) - (w) / (2.0 * syz**3)
                d2yi_dxkdxk = (3.0 * xi * xj**2)   / (4.0 * sjk**5) #d2y_dzdz = (3.0 * w * y**2)  / (4.0 * syz**5)

                d2y_dx2[i, i, j] = d2y_dx2[i, j, i] = d2yi_dxidxj
                d2y_dx2[i, i, k] = d2y_dx2[i, k, i] = d2yi_dxidxk
                d2y_dx2[i, j, j] = d2yi_dxjdxj
                d2y_dx2[i, j, k] = d2y_dx2[i, k, j] = d2yi_dxjdxk
                d2y_dx2[i, k, k] = d2yi_dxkdxk
        return d2y_dx2

    def _hess_rvs(self, y):
        rix, cix = self.row_diag_inds, self.col_diag_inds
        d2x_dy2 = np.zeros((self.hv_size,)*3)
        for i in range(self.hv_size):
            j, k = rix[i], cix[i]
            if j!=k:
               yi, yj, yk = y[i], y[j], y[k]
               d2xi_dyidyj = np.sqrt(yk) / (2.0 * np.sqrt(yj))
               d2xi_dyidyk = np.sqrt(yj) / (2.0 * np.sqrt(yk))
               d2xi_dyjdyj = (-yi * np.sqrt(yk)) / (4.0 * np.sqrt(yj)**3)
               d2xi_dyjdyk = yi / (4.0 * np.sqrt(yj * yk))
               d2xi_dykdyk = (-yi * np.sqrt(yj)) / (4.0 * np.sqrt(yk)**3)

               d2x_dy2[i, i, j] = d2x_dy2[i, j, i] = d2xi_dyidyj
               d2x_dy2[i, i, k] = d2x_dy2[i, k, i] = d2xi_dyidyk
               d2x_dy2[i, j, j] = d2xi_dyjdyj
               d2x_dy2[i, j, k] = d2x_dy2[i, k, j] = d2xi_dyjdyk
               d2x_dy2[i, k, k] = d2xi_dykdyk
        return d2x_dy2

class LogScale(object):

    def __init__(self, mat_size):
        self.mat_size = mat_size
        self.hv_size = mat_size_to_hv_size(mat_size)
        self.row_inds, self.col_inds = hv_indices((mat_size, mat_size))
        self.diag_inds, = np.where(self.row_inds==self.col_inds)
        self.tril_inds, = np.where(self.row_inds!=self.col_inds)

    def _fwd(self, x):
        ix = self.diag_inds
        y = x.copy()
        y[ix] = np.log(x[ix])
        return y

    def _rvs(self, y):
        ix = self.diag_inds
        x = y.copy()
        x[ix] = np.exp(y[ix])
        return x

    def _jac_fwd(self, x):
        ix = self.diag_inds
        dy_dx = np.eye(x.shape[-1])
        dy_dx[ix, ix] = 1 / x[ix]
        return dy_dx

    def _jac_rvs(self, y):
        ix = self.diag_inds
        dx_dy = np.eye(y.shape[-1])
        dx_dy[ix, ix] = np.exp(y[ix])
        return dx_dy

    def _hess_fwd(self, x):
        ix = self.diag_inds
        d2y_dx2 = np.zeros((x.shape[-1],)*3)
        d2y_dx2[ix, ix, ix] = -1 / x[ix]**2
        return d2y_dx2

    def _hess_rvs(self, y):
        ix = self.diag_inds
        d2x_dy2 = np.zeros((y.shape[-1],)*3)
        d2x_dy2[ix, ix, ix] = np.exp(y[ix])
        return d2x_dy2

class CorrCholesky(object):

    def __init__(self, mat_size):
        self.n = self.mat_size = mat_size
        self.m = self.vec_size = int((mat_size + 1) * mat_size / 2)
        self.I = np.eye(self.n)
        self.N = nmat(self.n)
        self.E = lmat(self.n)
        self.diag_inds = np.diag_indices(self.mat_size)
        self.lhv_size = mat_size_to_lhv_size(mat_size)
        self.row_inds, self.col_inds = lhv_indices((mat_size, mat_size))
        self.row_sort = np.argsort(self.row_inds)
        self.ind_parts = lhv_ind_parts(mat_size)
        self.row_norm_inds = [self.row_sort[a:b] for a, b in self.ind_parts]
        j, i = np.triu_indices(self.n)
        self.non_diag, = np.where(j!=i)
        self.dM, self.d2M = unit_matrices(mat_size)
        row_inds, col_inds = hv_indices((self.mat_size, self.mat_size))
        self.tril_inds, = np.where(row_inds!=col_inds)

    def _fwd(self, x):
        R = _invecl(x)
        L = np.linalg.cholesky(R)
        y = lower_half_vec(L)
        return y

    def _rvs(self, y):
        L = inv_lower_half_vec(y)
        L[self.diag_inds] = np.sqrt(1-np.linalg.norm(L, axis=-1)**2)
        R = np.dot(L, L.T)
        x = lower_half_vec(R)
        return x

    def _jac_fwd(self, x):
        M = _invecl(x)
        _, dL, _ = dchol(M, self.dM, self.d2M, order=1)
        dL = dL.reshape(np.product(dL.shape[:2]), dL.shape[2], order='F')
        dy_dx = self.E.dot(dL)[np.ix_(self.tril_inds, self.tril_inds)]
        return dy_dx

    def _jac_rvs(self, y):
        x = self._rvs(y)
        dx_dy = np.linalg.inv(self._jac_fwd(x))
        return dx_dy

    def _hess_fwd(self, x):
        M = _invecl(x)
        _, _, d2L = dchol(M, self.dM, self.d2M, order=2)
        k = (np.product(d2L.shape[:2]),)
        d2L = d2L.reshape(k+d2L.shape[2:], order='F')
        d2L = d2L[self.E.tocoo().col]
        d2y_dx2 = d2L[np.ix_(self.tril_inds, self.tril_inds, self.tril_inds)]
        return d2y_dx2

    def _hess_rvs(self, y):
        #x = self._rvs(y)
        #M = _invecl(x)
        #L, dL, d2L = dchol(M, self.dM, self.d2M, order=2)
        #k = (np.product(d2L.shape[:2]),)
        #d2L = d2L.reshape(k+d2L.shape[2:], order='F')
        #d2L = d2L[self.E.tocoo().col]
        #d2y_dx2 = d2L[np.ix_(self.tril_inds, self.tril_inds, self.tril_inds)]
        #dL = dL.reshape(np.product(dL.shape[:2]), dL.shape[2], order='F')
        #dL = self.E.dot(dL)[np.ix_(self.tril_inds, self.tril_inds)]
        #dx_dy = np.linalg.inv(dL)
        #d2x_dy2 = -np.einsum("ijk,kl->ijl", d2y_dx2, dx_dy)
        return jac_approx(self._jac_rvs, y.copy())

class OffDiagMask(object):

    def __init__(self, transform):
        self.mat_size = transform.mat_size
        self.hv_size = mat_size_to_hv_size(self.mat_size)
        self.row_inds, self.col_inds = hv_indices((self.mat_size, self.mat_size))
        self.diag_inds, = np.where(self.row_inds==self.col_inds)
        self.tril_inds, = np.where(self.row_inds!=self.col_inds)
        self.transform = transform

    def _fwd(self, x):
        y = x.copy()
        y[self.tril_inds] = self.transform._fwd(y[self.tril_inds])
        return y

    def _rvs(self, y):
        x = y.copy()
        x[self.tril_inds] = self.transform._rvs(x[self.tril_inds])
        return x

    def _jac_fwd(self, x):
        dy_dx = np.zeros((self.hv_size, self.hv_size))
        ii, ij = self.diag_inds, self.tril_inds
        dy_dx[np.ix_(ij, ij)] = self.transform._jac_fwd(x[ij].copy())
        dy_dx[np.ix_(ii, ii)] = np.eye(len(ii))
        return dy_dx

    def _jac_rvs(self, y):
        dx_dy = np.zeros((self.hv_size, self.hv_size))
        ii, ij = self.diag_inds, self.tril_inds
        dx_dy[np.ix_(ij, ij)] = self.transform._jac_rvs(y[ij].copy())
        dx_dy[np.ix_(ii, ii)] = np.eye(len(ii))
        return dx_dy

    def _hess_fwd(self, x):
        d2y_dx2 = np.zeros((self.hv_size, self.hv_size, self.hv_size))
        ij = self.tril_inds
        d2y_dx2[np.ix_(ij, ij, ij)] = self.transform._hess_fwd(x[ij].copy())
        return d2y_dx2

    def _hess_rvs(self, y):
        d2x_dy2 = np.zeros((self.hv_size, self.hv_size, self.hv_size))
        ij = self.tril_inds
        d2x_dy2[np.ix_(ij, ij, ij)] = self.transform._hess_rvs(y[ij].copy())
        return d2x_dy2

class CholeskyCov(object):

    def __init__(self, mat_size):
        self.covn_to_corr = CovCorr(mat_size)
        self.corr_to_chol = OffDiagMask(CorrCholesky(mat_size))
        self.chol_to_real = OffDiagMask(UnconstrainedCholeskyCorr(mat_size))
        self.vars_to_logs = LogScale(mat_size)

    def _fwd(self, x):
        y = self.covn_to_corr._fwd(x)
        z = self.corr_to_chol._fwd(y)
        w = self.chol_to_real._fwd(z)
        u = self.vars_to_logs._fwd(w)
        return u

    def _rvs(self, u):
        w = self.vars_to_logs._rvs(u)
        z = self.chol_to_real._rvs(w)
        y = self.corr_to_chol._rvs(z)
        x = self.covn_to_corr._rvs(y)
        return x

    def _jac_fwd(self, x):
        y = self.covn_to_corr._fwd(x)
        z = self.corr_to_chol._fwd(y)
        w = self.chol_to_real._fwd(z)
        #u = self.vars_to_logs._fwd(w)

        dy_dx = self.covn_to_corr._jac_fwd(x)
        dz_dy = self.corr_to_chol._jac_fwd(y)
        dw_dz = self.chol_to_real._jac_fwd(z)
        du_dw = self.vars_to_logs._jac_fwd(w)
        dw_dx = dw_dz.dot(dz_dy).dot(dy_dx)
        du_dx = du_dw.dot(dw_dx)
        return du_dx

    def _jac_rvs(self, u):
        w = self.vars_to_logs._rvs(u)
        z = self.chol_to_real._rvs(w)
        y = self.corr_to_chol._rvs(z)

        dx_dy = self.covn_to_corr._jac_rvs(y)
        dy_dz = self.corr_to_chol._jac_rvs(z)
        dz_dw = self.chol_to_real._jac_rvs(w)
        dw_du = self.vars_to_logs._jac_rvs(u)
        dx_dw = dx_dy.dot(dy_dz).dot(dz_dw)
        dx_du = dx_dw.dot(dw_du)
        return dx_du

    def _hess_fwd(self, x):
        y = self.covn_to_corr._fwd(x)
        z = self.corr_to_chol._fwd(y)
        w = self.chol_to_real._fwd(z)
        dy_dx = self.covn_to_corr._jac_fwd(x)
        dz_dy = self.corr_to_chol._jac_fwd(y)
        dw_dz = self.chol_to_real._jac_fwd(z)
        du_dw = self.vars_to_logs._jac_fwd(w)

        dz_dx = dz_dy.dot(dy_dx)
        dw_dx = dw_dz.dot(dz_dy.dot(dy_dx))
        d2y_dx2 = self.covn_to_corr._hess_fwd(x)
        d2z_dy2 = self.corr_to_chol._hess_fwd(y)
        d2w_dz2 = self.chol_to_real._hess_fwd(z)
        d2u_dw2 = self.vars_to_logs._hess_fwd(w)

        d2z_dx2 = _hess_chain_rule(d2z_dy2, dy_dx, dz_dy, d2y_dx2)
        d2w_dx2 = _hess_chain_rule(d2w_dz2, dz_dx, dw_dz, d2z_dx2)
        d2u_dx2 = _hess_chain_rule(d2u_dw2, dw_dx, du_dw, d2w_dx2)

        # m = x.shape[-1]
        # d2z_dx2 = np.zeros((m,)*3)
        # for i in range(m):
        #     d2z_dx2[i] += dy_dx.T.dot(d2z_dy2[i]).dot(dy_dx)
        #     for j in range(m):
        #         d2z_dx2[i] += dz_dy[i, j] * d2y_dx2[j]
        # d2w_dx2 = np.zeros((m,)*3)
        # for i in range(m):
        #     d2w_dx2[i] += dz_dx.T.dot(d2w_dz2[i]).dot(dz_dx)
        #     for j in range(m):
        #         d2w_dx2[i] += dw_dz[i, j] * d2z_dx2[j]
        # d2u_dx2 = np.zeros((m,)*3)
        # for i in range(m):
        #     d2u_dx2[i] += dw_dx.T.dot(d2u_dw2[i]).dot(dw_dx)
        #     for j in range(m):
        #         d2u_dx2[i] += du_dw[i, j] * d2w_dx2[j]
        return d2u_dx2

max_stars_repo_name: lukepinkel/pystatsm

max_issues_repo_name: lukepinkel/pystatsm

max_forks_repo_name: lukepinkel/pystatsm

reason: import numpy


text:
import re
from collections import defaultdict, Counter
from numbers import Number
from typing import Optional, List, Dict
import numpy as np

import torch
from allennlp.common import FromParams, Registrable
from dataclasses import dataclass, replace
from pycocoevalcap.bleu.bleu import Bleu
from pycocoevalcap.cider.cider import Cider
import third_party.detection_metrics.lib.Evaluator as det_evaluator

from gpv2.data.dataset import VqaExample, ClsExample, WebQaExample, LocalizationExample, \
  CaptioningExample
from gpv2.data.gpv_datasets import COCO_CATEGORIES
from gpv2.data.synonyms import SYNONYMS
from gpv2.model.model import GPVExampleOutput
from gpv2.train.vqa2_eval_data import punct, periodStrip, commaStrip, manualMap, articles, \
  contractions
from gpv2.utils import py_utils
from gpv2.utils.image_utils import get_image_size
from gpv2.utils.quiet_ptbtokenizer import QuitePTBTokenizer

def vqa_score(answer, ground_truth_answer_counts):
  gt_answers = {k.lower(): v for k, v in ground_truth_answer_counts.items()}
  return min(gt_answers.get(answer, 0) / 3, 1)

@dataclass(frozen=True)
class ResultKey(FromParams):
  """Key for a result from a model"""

  metric_name: str
  subset_name: Optional[str] = None
  dataset_name: Optional[str] = None

  def __str__(self):
    out = [self.dataset_name, self.subset_name, self.metric_name]
    return "/".join(x for x in out if x is not None)

  def __repr__(self):
    return str(self)

class Evaluator(Registrable):
  """Computes evaluations metrics"""

  def evaluate(
      self, examples: List, predictions: Dict[str, GPVExampleOutput],
      allow_partial=False, subset_mapping=None
  ) -> Dict[ResultKey, Number]:
    """Computes corpus wide metrics

    :param examples: List of source examples
    :param predictions: example key -> model output
    :param allow_partial: Allow the predictions to only cover a subset of `examples`,
                          in which only those predictions should be evaluated
    :param subset_mapping: Function that maps example -> list of strings, names of the subsets that
                           example is part of
    """
    raise NotImplementedError()

class PerExampleEvaluator(Evaluator):
  """Computes per-examples evaluations metrics"""

  def evaluate_examples(self, examples: List, predictions: Dict[str, GPVExampleOutput])-> List[Dict[str, Number]]:
    raise NotImplementedError()

  def evaluate(
      self,
      examples: List,
      predictions: Dict[str, GPVExampleOutput],
      allow_partial=False,
      mean=True,
      subset_mapping=None
  ) -> Dict[ResultKey, Number]:
    examples_with_predictions = [x for x in examples if x.get_gpv_id() in predictions]
    if not allow_partial and (len(examples) != len(examples_with_predictions)):
      raise ValueError(f"Only {len(examples_with_predictions)}/{len(examples)} "
                       f"of examples have predictions")
    examples = examples_with_predictions
    per_example_scores = self.evaluate_examples(examples, predictions)
    per_metric_scores = py_utils.transpose_list_of_dicts(per_example_scores)

    subsets = defaultdict(list)
    all_ids = [x.get_gpv_id() for x in examples]

    id_to_ix = {k: i for i, k in enumerate(all_ids)}
    subsets[None] = list(range(len(all_ids)))

    if subset_mapping is not None:
      for example in examples:
        example_id = id_to_ix[example.get_gpv_id()]
        for subset in subset_mapping(example):
          subsets[subset].append(example_id)

    out = {}
    for metric_name, score in per_metric_scores.items():
      score = np.array(score)
      for subset_name, ixs in subsets.items():
        if mean:
          out[ResultKey(metric_name, subset_name)] = float(np.mean(score[ixs]))
        else:
          out[ResultKey(metric_name, subset_name)] = (float(np.sum(score[ixs])), len(ixs))

    return out

@Evaluator.register("vqa-evaluator")
class VqaEvaluator(PerExampleEvaluator):

  def evaluate_examples(self, examples: List[VqaExample],
                        predictions: Dict[str, GPVExampleOutput], add_scores=False):
    out = []
    for example in examples:
      answer = predictions[example.gpv_id].text[0]
      score = vqa_score(answer.lower(), example.answers)
      out.append(dict(score=score))
    return out

@Evaluator.register("cls-evaluator")
class ClsEvaluator(PerExampleEvaluator):

  def evaluate_examples(self, examples: List[ClsExample], predictions: Dict[str, GPVExampleOutput]):
    out = []
    for example in examples:
      answer = predictions[example.gpv_id].text[0].lower()
      gt_answer = SYNONYMS[example.category]
      out.append(dict(accuracy=answer in gt_answer))
    return out

@Evaluator.register("webqa-evaluator")
class WebQaEvaluator(PerExampleEvaluator):
  def evaluate_examples(self, examples: List[WebQaExample], predictions: Dict[str, GPVExampleOutput]):
    out = []
    for example in examples:
      answer = predictions[example.get_gpv_id()].text[0].lower()
      gt_answer = SYNONYMS[example.answer] if example.answer in SYNONYMS else [example.answer]
      out.append(dict(accuracy=answer in gt_answer))
    return out

@Evaluator.register("dce-cls")
class DceClsEvaluator(PerExampleEvaluator):

  def __init__(self, top_k: Optional[List[int]]=(5,)):
    self.top_k = top_k

  def evaluate_examples(self, examples: List[ClsExample], predictions: Dict[str, GPVExampleOutput]):
    out = []
    for example in examples:
      answers = [x.lower() for x in predictions[example.get_gpv_id()].text]
      gt = [example.category]
      vals = {"accuracy": answers[0] in gt}
      if self.top_k is not None:
        for k in self.top_k:
          assert len(answers) >= k
          vals[f"top{k}-acc"] = any(a in gt for a in answers[:k])
      out.append(vals)
    return out

def compute_vqa_accuracy(
    gt_answers: List[str],
    pred_answers: List[str]) -> List[float]:
  ngt_answers = [preprocess_answer(ans) for ans in gt_answers]
  topk_npred_answers = [preprocess_answer(ans) for ans in pred_answers]
  gt_consensus = Counter(ngt_answers)
  return [vqa_accuracy(ans, gt_consensus) for ans in topk_npred_answers]

def vqa_accuracy(npred_answer: str, gt_consensus: Counter):
  return min(gt_consensus[npred_answer]/3,1)

def processPunctuation(inText):
  outText = inText
  for p in punct:
    if (p + ' ' in inText or ' ' + p in inText) or (re.search(commaStrip, inText) != None):
      outText = outText.replace(p, '')
    else:
      outText = outText.replace(p, ' ')
  outText = periodStrip.sub("",outText,re.UNICODE)
  return outText

def processDigitArticle(inText):
  outText = []
  tempText = inText.lower().split()
  for word in tempText:
    word = manualMap.setdefault(word, word)
    if word not in articles:
      outText.append(word)
    else:
      pass
  for wordId, word in enumerate(outText):
    if word in contractions:
      outText[wordId] = contractions[word]
  outText = ' '.join(outText)
  return outText

def preprocess_answer(ans):
  ans = ans.replace('\n', ' ')
  ans = ans.replace('\t',' ')
  ans = ans.lower().strip()
  return processDigitArticle(processPunctuation(ans))

@Evaluator.register("opensce-vqa")
class DceVqaEvaluator(PerExampleEvaluator):

  def __init__(self, top_k: Optional[List[int]]=(5,)):
    self.top_k = top_k
    if top_k is not None:
      assert all(x > 0 for x in top_k)

  def evaluate_examples(self, examples: List[VqaExample], predictions: Dict[str, GPVExampleOutput]):
    out = []
    for example in examples:
      max_k = 1 if self.top_k is None else max(self.top_k)
      answers = predictions[example.get_gpv_id()].text[:max_k]
      gt = example.answers
      scores = compute_vqa_accuracy(gt, answers)
      vals = dict(acc=scores[0])
      if self.top_k:
        for k in self.top_k:
          if k > len(scores):
            raise ValueError(f"Cannot evaluate top-{k}, but only have top-{len(scores)} predictions")
          vals[f"top{k}-acc"] = max(scores[:k])
      out.append(vals)
    return out

@Evaluator.register("localization-evaluator")
@Evaluator.register("detect-evaluator")
class LocalizationEvaluator(PerExampleEvaluator):

  def __init__(self, iou_thresh=0.5):
    self.iou_thresh = iou_thresh

  def evaluate_examples(self, examples: List[LocalizationExample], predictions: Dict[str, GPVExampleOutput],
                        return_pr=False):
    eval_engine = det_evaluator.Evaluator()
    out = []

    for i, ex in enumerate(examples):
      pred = predictions[ex.gpv_id]
      scores = pred.relevance
      pred_boxes = pred.boxes.copy()
      gt_boxes = np.array(ex.bboxes)

      # Convert cx cy, w, h -> x1, y1, w, h
      pred_boxes[:, 0] = pred_boxes[:, 0] - 0.5 * pred_boxes[:, 2]
      pred_boxes[:, 1] = pred_boxes[:, 1] - 0.5 * pred_boxes[:, 3]

      B = pred_boxes.shape[0]
      all_boxes = det_evaluator.BoundingBoxes()
      W, H = get_image_size(ex.image_id)

      for b in range(B):
        x, y, w, h = pred_boxes[b]
        all_boxes.addBoundingBox(det_evaluator.BoundingBox(
          imageName=ex.image_id,
          classId=ex.category,
          x=x,
          y=y,
          w=w,
          h=h,
          typeCoordinates=det_evaluator.CoordinatesType.Relative,
          imgSize=(W, H),
          bbType=det_evaluator.BBType.Detected,
          classConfidence=scores[b],
          format=det_evaluator.BBFormat.XYWH))

      normalized_gt = all(all(val <= 1.0 for val in b) for b in gt_boxes)
      if not normalized_gt:
        # convert to relative coordinates
        # TODO its a bit of hack to check this by looking coordinates > 1.0
        # but we need this check atm since DCE stores relative scaling
        # coco uses absolute
        W, H = get_image_size(ex.image_id)
        gt_boxes[:, 0] = gt_boxes[:, 0] / W
        gt_boxes[:, 1] = gt_boxes[:, 1] / H
        gt_boxes[:, 2] = gt_boxes[:, 2] / W
        gt_boxes[:, 3] = gt_boxes[:, 3] / H

      B = gt_boxes.shape[0]
      for b in range(B):
        x, y, w, h = gt_boxes[b]

        all_boxes.addBoundingBox(det_evaluator.BoundingBox(
          imageName=ex.image_id,
          classId=ex.category,
          x=x,
          y=y,
          w=w,
          h=h,
          typeCoordinates=det_evaluator.CoordinatesType.Relative,
          imgSize=(W, H),
          bbType=det_evaluator.BBType.GroundTruth,
          format=det_evaluator.BBFormat.XYWH))

      det_metrics = eval_engine.GetPascalVOCMetrics(all_boxes, self.iou_thresh)
      if return_pr:
        out.append(det_metrics[0])
      else:
        out.append({"AP": det_metrics[0]['AP']})

    return out

def get_per_caption_data(examples: List[CaptioningExample], predictions):
  # In per-caption evaluation the model makes one prediction for each ground truth
  # caption, each of which it still evaluated against all the captions,
  caption_examples = []
  caption_predictions = {}
  for ex in examples:
    pred = predictions[ex.gpv_id]
    for cap in ex.captions:
      caption_examples.append(CaptioningExample(cap.gpv_id, ex.image_id, ex.captions))
      caption_predictions[cap.gpv_id] = pred
  return caption_examples, caption_predictions

@Evaluator.register("cap-evaluator")
class CaptionEvaluator(Evaluator):
  def __init__(self, cider=True, bleu=4, per_caption=False):
    self.cider = cider
    self.bleu = bleu
    self.per_caption = per_caption
    scorers = {}
    if cider:
      # from exp.ours.eval.fast_cider import FastCider
      scorers["cider"] = Cider()
    if bleu:
      scorers["bleu"] = Bleu(bleu)
    self.scorers = scorers
    self.tokenizer = QuitePTBTokenizer()

  def evaluate(
      self,
      examples: List,
      predictions: Dict[str, GPVExampleOutput],
      allow_partial=False,
      subset_mapping=None,
  ):
    examples_with_predictions = [x for x in examples if x.get_gpv_id() in predictions]
    if not allow_partial and (len(examples) != len(examples_with_predictions)):
      raise ValueError(f"Only {len(examples_with_predictions)}/{len(examples)} "
                       f"of examples have predictions")
    examples = examples_with_predictions

    if self.per_caption:
      examples, predictions = get_per_caption_data(examples, predictions)

    subsets = defaultdict(list)
    subsets[None] = examples
    if subset_mapping is not None:
      for example in examples:
        example_subsets = subset_mapping(example)
        for subset in example_subsets:
          subsets[subset].append(example)

    out = {}
    for subset_name, examples in subsets.items():
      all_scores = self._get_scores(examples, predictions)

      results = {}
      for name, scorer in self.scorers.items():
        corpus_scores, _ = all_scores[name]
        if isinstance(scorer, Cider):
          results["cider"] = corpus_scores
        elif isinstance(scorer, Bleu):
          scores, _ = all_scores[name]
          for i, score in enumerate(corpus_scores):
            results[f"bleu{i+1}"] = score
      if subset_name is not None:
        results["n"] = len(examples)
      out.update({ResultKey(metric_name=k, subset_name=subset_name): v for k, v in results.items()})
    return out

  def evaluate_examples(self, examples: List[CaptioningExample], predictions: Dict[str, GPVExampleOutput]):
    all_scores = self._get_scores(examples, predictions)

    per_examples_scores = [{} for _ in examples]
    for name, scorer in self.scorers.items():
      score, scores = all_scores[name]
      if isinstance(scorer, Cider):
        for score, ex_scores in zip(scores, per_examples_scores):
          ex_scores["cider"] = score
      elif isinstance(scorer, Bleu):
        scores = py_utils.transpose_lists(scores)
        for score, ex_scores in zip(scores, per_examples_scores):
          for i, s in enumerate(score):
            ex_scores[f"bleu{i+1}"] = s

    return per_examples_scores

  def _get_scores(self,  examples: List[CaptioningExample], predictions: Dict[str, GPVExampleOutput]):
    gts = {}
    res = {}
    for ix, instance in enumerate(examples):
      key = instance.get_gpv_id()
      assert key not in res
      res[key] = [predictions[instance.get_gpv_id()].text[0]]
      gts[key] = [x.caption.lower() for x in instance.captions]

    res = self.tokenizer.tokenize(res)
    gts = self.tokenizer.tokenize(gts)

    scores = {}
    for name, scorer in self.scorers.items():
      if isinstance(scorer, Bleu):
        scores[name] = scorer.compute_score(gts, res, verbose=0)
      else:
        scores[name] = scorer.compute_score(gts, res)
    return scores