text:
#import .config as config
from __future__ import absolute_import

import sys
import os

try:
    from dotenv import find_dotenv, load_dotenv
except:
    pass

import argparse
from tqdm import tqdm
import json

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

from data import load_data, data_utils
from models import transformer_models, activations, layers, losses, scorers

from .trainer_utils import set_seed, _has_apex, _torch_lightning_available, _has_wandb, _torch_gpu_available, _num_gpus, _torch_tpu_available

import pandas as pd
import numpy as np

import torch
import torch.nn as nn
from torchcontrib.optim import SWA
from torch.optim import Adam, SGD
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau, CyclicLR, \
                                     CosineAnnealingWarmRestarts

from transformers import AdamW, get_linear_schedule_with_warmup

if _has_apex:
    #from torch.cuda import amp
    from apex import amp

if _torch_tpu_available:
    import torch_xla
    import torch_xla.core.xla_model as xm
    import torch_xla.distributed.xla_multiprocessing as xmp

if _has_wandb:
    import wandb
    try:
        load_dotenv(find_dotenv())
        wandb.login(key=os.environ['WANDB_API_KEY'])
    except:
        _has_wandb = False

if _torch_lightning_available:
    import pytorch_lightning as pl
    from pytorch_lightning import Trainer, seed_everything
    from pytorch_lightning.loggers import WandbLogger
    from pytorch_lightning.metrics.metric import NumpyMetric
    from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, Callback

def test_fn(data_loader, model, device, final_activation=None):
    model.eval()
    fin_targets = []
    fin_outputs = []
    with torch.no_grad():
        for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
            ids = d["ids"]
            token_type_ids = d["token_type_ids"]
            mask = d["mask"]

            ids = ids.to(device, dtype=torch.long)
            token_type_ids = token_type_ids.to(device, dtype=torch.long)
            mask = mask.to(device, dtype=torch.long)

            outputs = model(ids=ids, mask=mask, token_type_ids=token_type_ids)
            if final_activation == 'sigmoid':
                fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
            else:
                fin_outputs.extend(outputs.cpu().detach().numpy().tolist())

    return fin_outputs

def test_fn_multi(data_loader1, data_loader2, model, device, final_activation=None):
    model.eval()
    fin_targets = []
    fin_outputs = []
    with torch.no_grad():
        for bi, (d1, d2) in tqdm(enumerate(zip(data_loader1,data_loader2)), total=len(data_loader1)):
            ids1 = d1["ids"]
            token_type_ids1= d1["token_type_ids"]
            mask1 = d1["mask"]

            ids2 = d2["ids"]
            token_type_ids2= d2["token_type_ids"]
            mask2 = d2["mask"]

            ids1 = ids1.to(device, dtype=torch.long)
            token_type_ids1 = token_type_ids1.to(device, dtype=torch.long)
            mask1 = mask1.to(device, dtype=torch.long)

            ids2 = ids2.to(device, dtype=torch.long)
            token_type_ids2 = token_type_ids2.to(device, dtype=torch.long)
            mask2 = mask2.to(device, dtype=torch.long)

            outputs = model(ids1,ids2, mask1,mask2,token_type_ids1,token_type_ids2)
            if final_activation == 'sigmoid':
                fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
            else:
                fin_outputs.extend(outputs.cpu().detach().numpy().tolist())

    return fin_outputs

def test_fn_qa(data_loader, model, device, final_activation=None):
    model.eval()
    fin_outputs_start = []
    fin_outputs_end = []

    with torch.no_grad():
        for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
            ids = d["ids"]
            token_type_ids = d["token_type_ids"]
            mask = d["mask"]

            ids = ids.to(device, dtype=torch.long)
            token_type_ids = token_type_ids.to(device, dtype=torch.long)
            mask = mask.to(device, dtype=torch.long)

            outputs_start, outputs_end = model(ids=ids, mask=mask, token_type_ids=token_type_ids)
            if final_activation == 'sigmoid':
                fin_outputs_start.extend(torch.sigmoid(outputs_start).cpu().detach().numpy().tolist())
                fin_outputs_end.extend(torch.sigmoid(outputs_end).cpu().detach().numpy().tolist())
            else:
                fin_outputs_start.extend(outputs_start.cpu().detach().numpy().tolist())
                fin_outputs_end.extend(outputs_end.cpu().detach().numpy().tolist())

    return fin_outputs_start, fin_outputs_end

def test_pl_trainer(data_loader, pltrainer, final_activation=None):
    fin_targets = []
    fin_outputs = []
    with torch.no_grad():
        for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
            outputs = pltrainer(d)
            if final_activation == 'sigmoid':
                fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
            else:
                fin_outputs.extend(outputs.cpu().detach().numpy().tolist())

    return fin_outputs

class BasicTrainer:

    def __init__(self, model, train_data_loader, val_data_loader, device, model_description, final_activation=None, test_data_loader=None):
        self.model = model
        self.train_data_loader = train_data_loader
        self.val_data_loader = val_data_loader
        self.device = device
        self.description = model_description
        self.test_data_loader = test_data_loader
        self.final_activation = final_activation

        self.model.to(self.device)

        self.print_stats()

    def print_stats(self):
        print ("[LOG] Total number of parameters to learn {}".format(sum(p.numel() for p in self.model.parameters() \
                                                                 if p.requires_grad)))

    def train_fn(self, data_loader):
        self.model.train()

        total_loss = 0
        pbar = tqdm(enumerate(data_loader), total=len(data_loader))
        for bi, d in pbar:
            ids = d["ids"]
            token_type_ids = d["token_type_ids"]
            mask = d["mask"]
            targets = d["targets"]

            ids = ids.to(self.device)
            token_type_ids = token_type_ids.to(self.device)
            mask = mask.to(self.device)
            targets = targets.to(self.device)

            self.optimizer.zero_grad()

            if self.is_amp:
                #with amp.autocast():
                outputs = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
                if self.loss_fn.__class__.__name__ == 'masked_CELoss':
                    loss = self.loss_fn(outputs, targets, mask)
                else:
                    loss = self.loss_fn(outputs, targets)
            else:
                outputs = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
                if self.loss_fn.__class__.__name__ == 'masked_CELoss':
                    loss = self.loss_fn(outputs, targets, mask)
                else:
                    loss = self.loss_fn(outputs, targets)

            #nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

            if self.scaler:
                #self.scaler.scale(loss.backward())
                #self.scaler.step(self.optimizer.step())
                #self.scaler.update()
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
                self.optimizer.step()
                self.scheduler.step()

            else:
                loss.backward()
                if self.num_tpus == 0:
                    self.optimizer.step()
                else:
                    xm.optimizer_step(self.optimizer, barrier=True)

                self.scheduler.step()

            total_loss += loss
            with np.printoptions(precision=3):
                v = round(loss.cpu().detach().numpy().item(), 3)
                pbar.set_description("Current training Loss {}".format(v))

        return total_loss

    def eval_fn(self, data_loader):
        self.model.eval()
        fin_targets = []
        fin_outputs = []

        total_loss = 0
        with torch.no_grad():
            pbar = tqdm(enumerate(data_loader), total=len(data_loader))
            for bi, d in pbar:
                ids = d["ids"]
                token_type_ids = d["token_type_ids"]
                mask = d["mask"]
                targets = d["targets"]

                ids = ids.to(self.device)
                token_type_ids = token_type_ids.to(self.device)
                mask = mask.to(self.device)
                targets = targets.to(self.device)

                outputs = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
                if self.loss_fn.__class__.__name__ == 'masked_CELoss':
                    loss = self.loss_fn(outputs, targets, mask)
                else:
                    loss = self.loss_fn(outputs, targets)

                total_loss += loss

                with np.printoptions(precision=3):
                    v = round(loss.cpu().detach().numpy().item(), 3)
                    pbar.set_description("Current eval Loss {}".format(v))

                fin_targets.extend(targets.cpu().detach().numpy().tolist())
                if self.final_activation == 'sigmoid':
                    fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
                else:
                    fin_outputs.extend(outputs.cpu().detach().numpy().tolist())

        return fin_outputs, total_loss, fin_targets

    def train(self, epochs, lr, scorer, loss_fn, optimizer, scheduler, MODEL_PATH, num_gpus=0, num_tpus=0, max_grad_norm=1, early_stopping_rounds=3, \
                    snapshot_ensemble=False, is_amp=False, use_wandb=False, seed=42):

        self.scorer = scorer
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.num_gpus = num_gpus
        self.num_tpus = num_tpus

        set_seed(seed)

        self.args = {'description': self.description, 'loss': loss_fn.__class__.__name__, 'epochs': epochs, 'learning_rate': lr, \
                    'optimizer': self.optimizer.__class__.__name__, 'scheduler': self.scheduler.__class__.__name__}

        if use_wandb and not _torch_tpu_available:
            wandb.init(project="Project",config=self.args)
            wandb.watch(self.model)
        else:
            with open(os.path.join(MODEL_PATH, 'parameters.json'),'w') as outfile:
                json.dump(self.args, outfile)

        if _num_gpus > 1 and (num_gpus != 0 and num_gpus != 1):
            self.model = nn.DataParallel(self.model)

        if is_amp and _has_apex:
            #scaler = amp.GradScaler()
            scaler = 1
        else:
            scaler = None

        if os.path.exists(os.path.join(MODEL_PATH,'model.bin')):
            try:
                try:
                    #self.model.load_state_dict(torch.load(os.path.join(MODEL_PATH,'model.bin')))
                    self.model.load_state_dict(torch.load(os.path.join(MODEL_PATH,'model.bin')))
                except:
                    self.model = torch.load(os.path.join(MODEL_PATH,'model.bin'))

                #print (self.model)
                print ("Loaded model from previous checkpoint")
            except:
                pass

        self.is_amp = is_amp
        self.max_grad_norm = max_grad_norm
        self.scaler = scaler

        best_metric = 0
        bad_epochs = 0
        if snapshot_ensemble:
            test_outputs = []
            val_outputs = []

        stats = {}
        for epoch in range(epochs):
            if bad_epochs < early_stopping_rounds:
                train_loss = self.train_fn(self.train_data_loader)

                print ("Running evaluation on whole training data")
                train_out, train_loss, train_targets = self.eval_fn(self.train_data_loader)
                print ("Running evaluation on validation data")
                val_out, val_loss, val_targets = self.eval_fn(self.val_data_loader)

                train_loss = train_loss/len(self.train_data_loader)
                val_loss = val_loss/len(self.val_data_loader)

                #train_out = np.round(train_out)

                if snapshot_ensemble:
                    val_outputs += [np.expand_dims(np.array(val_out),axis=-1)]
                    val_out = np.concatenate(val_outputs,axis=-1).mean(axis=-1)
                #else:
                #    val_out = np.round(val_out) #np.array(outputs) >= 0.5

                train_metric = scorer(train_targets, train_out)
                val_metric = scorer(val_targets,val_out)

                with np.printoptions(precision=3):
                    print("Train loss = {} Train metric = {} Val loss = {} Val metric = {}".format(round(train_loss.detach().cpu().numpy().item(), 3),round(train_metric,3),\
                        round(val_loss.detach().cpu().numpy().item(), 3),round(val_metric, 3)))

                if val_metric > best_metric:
                    #torch.save(self.model.state_dict(), os.path.join(MODEL_PATH,'model.bin'))
                    torch.save(self.model, os.path.join(MODEL_PATH,'model.bin'))
                    best_metric = val_metric
                    bad_epochs = 0

                else:
                    bad_epochs += 1

                if snapshot_ensemble and self.test_data_loader:
                    test_out = test_fn(self.test_data_loader, self.model, self.device, self.final_activation)

                    test_outputs += [np.expand_dims(np.array(test_out),axis=-1)]

            stats.update({"epoch_{}".format(epoch): {"train_loss": round(train_loss.detach().cpu().numpy().item(), 3), "train_metric": round(train_metric,3), \
                "val_loss": round(val_loss.detach().cpu().numpy().item(), 3),  "val_metric": round(val_metric,3)}})

            if use_wandb:
                wandb.log({"train_loss": train_loss, "val_loss": val_loss, "train_metric": train_metric, "val_metric": val_metric})

        if self.test_data_loader:
            if snapshot_ensemble:
                self.test_output = np.concatenate(test_outputs,axis=-1).mean(axis=-1)
            else:
                self.test_output = test_fn(self.test_data_loader, self.model, self.device, self.final_activation)
        else:
            self.test_output = []

        #with np.printoptions(precision=3):
        #    print (stats)

        if use_wandb == False:
            with open(os.path.join(MODEL_PATH,'all_stats.json'), 'w') as outfile:
                json.dump(stats, outfile)

            with open(os.path.join(MODEL_PATH,'final_stats.json'), 'w') as outfile:
                d = {"epoch": epoch, "train_loss": round(train_loss.detach().cpu().numpy().item(), 3), "train_metric": round(train_metric,3), \
                "val_loss": round(val_loss.detach().cpu().numpy().item(), 3),  "val_metric": round(val_metric,3)}
                json.dump(d, outfile)

class BasicMultiTrainer:

    def __init__(self, model, train_data_loader1, train_data_loader2, val_data_loader1, val_data_loader2, device, model_description, final_activation=None, \
                 test_data_loader1=None, test_data_loader2=None):
        self.model = model
        self.train_data_loader1 = train_data_loader1
        self.train_data_loader2 = train_data_loader2
        self.val_data_loader1 = val_data_loader1
        self.val_data_loader2 = val_data_loader2
        self.device = device
        self.description = model_description
        self.final_activation = final_activation
        self.test_data_loader1 = test_data_loader1
        self.test_data_loader2 = test_data_loader2

        self.model.to(self.device)

        self.print_stats()

    def print_stats(self):
        print ("[LOG] Total number of parameters to learn {}".format(sum(p.numel() for p in self.model.parameters() \
                                                                 if p.requires_grad)))

    def train_fn(self, data_loader1, data_loader2):
        self.model.train()

        total_loss = 0
        pbar = tqdm(enumerate(zip(data_loader1,data_loader2)), total=len(data_loader1))
        for bi, (d1, d2) in pbar:
            ids1 = d1["ids"]
            token_type_ids1= d1["token_type_ids"]
            mask1 = d1["mask"]

            ids2 = d2["ids"]
            token_type_ids2= d2["token_type_ids"]
            mask2 = d2["mask"]

            targets = d1["targets"]

            ids1 = ids1.to(self.device)
            token_type_ids1 = token_type_ids1.to(self.device)
            mask1 = mask1.to(self.device)

            ids2 = ids2.to(self.device)
            token_type_ids2 = token_type_ids2.to(self.device)
            mask2 = mask2.to(self.device)

            targets = targets.to(self.device)

            self.optimizer.zero_grad()

            if self.is_amp:
                #with amp.autocast():
                outputs = self.model(ids1,ids2, mask1,mask2,token_type_ids1,token_type_ids2)
                if self.loss_fn.__class__.__name__ == 'masked_CELoss':
                    loss = self.loss_fn(outputs, targets, mask1)
                else:
                    loss = self.loss_fn(outputs, targets)

            else:
                outputs = self.model(ids1,ids2, mask1,mask2,token_type_ids1,token_type_ids2)
                if self.loss_fn.__class__.__name__ == 'masked_CELoss':
                    loss = self.loss_fn(outputs, targets, mask1)
                else:
                    loss = self.loss_fn(outputs, targets)

            #nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

            if self.scaler:
                #self.scaler.scale(loss.backward())
                #self.scaler.step(self.optimizer.step())
                #self.scaler.update()

                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()

                self.optimizer.step()
                self.scheduler.step()

            else:
                loss.backward()
                if self.num_tpus == 0:
                    self.optimizer.step()
                else:
                    xm.optimizer_step(self.optimizer, barrier=True)

                self.scheduler.step()

            total_loss += loss

            with np.printoptions(precision=3):
                v = round(loss.cpu().detach().numpy().item(), 3)
                pbar.set_description("Current training Loss {}".format(v))

        return total_loss

    def eval_fn(self, data_loader1, data_loader2):
        self.model.eval()
        fin_targets = []
        fin_outputs = []

        total_loss = 0
        with torch.no_grad():
            pbar = tqdm(enumerate(zip(data_loader1,data_loader2)), total=len(data_loader1))
            for bi, (d1, d2) in pbar:
                ids1 = d1["ids"]
                token_type_ids1= d1["token_type_ids"]
                mask1 = d1["mask"]

                ids2 = d2["ids"]
                token_type_ids2= d2["token_type_ids"]
                mask2 = d2["mask"]

                targets = d1["targets"]

                ids1 = ids1.to(self.device)
                token_type_ids1 = token_type_ids1.to(self.device)
                mask1 = mask1.to(self.device)

                ids2 = ids2.to(self.device)
                token_type_ids2 = token_type_ids2.to(self.device)
                mask2 = mask2.to(self.device)

                targets = targets.to(self.device)

                outputs = self.model(ids1,ids2, mask1,mask2,token_type_ids1,token_type_ids2)
                if self.loss_fn.__class__.__name__ == 'masked_CELoss':
                    loss = self.loss_fn(outputs, targets, mask1)
                else:
                    loss = self.loss_fn(outputs, targets)

                total_loss += loss

                with np.printoptions(precision=3):
                    v = round(loss.cpu().detach().numpy().item(), 3)
                    pbar.set_description("Current eval Loss {}".format(v))

                fin_targets.extend(targets.cpu().detach().numpy().tolist())
                if self.final_activation == 'sigmoid':
                    fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
                else:
                    fin_outputs.extend(outputs.cpu().detach().numpy().tolist())

        return fin_outputs, total_loss, fin_targets

    def train(self, epochs, lr, scorer, loss_fn, optimizer, scheduler, MODEL_PATH, num_gpus=0, num_tpus=0, max_grad_norm=1, early_stopping_rounds=3, \
                    snapshot_ensemble=False, is_amp=False, use_wandb=False, seed=42):

        self.scorer = scorer
        self.loss_fn = loss_fn #losses.get_loss(loss_fn)
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.num_gpus = num_gpus
        self.num_tpus = num_tpus

        set_seed(seed)

        self.args = {'description': self.description, 'loss': loss_fn.__class__.__name__, 'epochs': epochs, 'learning_rate': lr, \
                    'optimizer': self.optimizer.__class__.__name__, 'scheduler': self.scheduler.__class__.__name__}

        if use_wandb and not _torch_tpu_available:
            wandb.init(project="Project",config=self.args)
            wandb.watch(self.model)
        else:
            with open(os.path.join(MODEL_PATH, 'parameters.json'),'w') as outfile:
                json.dump(self.args, outfile)

        if _num_gpus > 1 and (num_gpus != 0 and num_gpus != 1):
            self.model = nn.DataParallel(self.model)

        if is_amp and _has_apex:
            #scaler = amp.GradScaler()
            scaler = 1
        else:
            scaler = None

        if os.path.exists(os.path.join(MODEL_PATH,'model.bin')):
            try:
                try:
                    #self.model.load_state_dict(torch.load(os.path.join(MODEL_PATH,'model.bin')))
                    self.model.load_state_dict(torch.load(os.path.join(MODEL_PATH,'model.bin')))
                except:
                    self.model = torch.load(os.path.join(MODEL_PATH,'model.bin'))

                #print (self.model)
                print ("Loaded model from previous checkpoint")
            except:
                pass

        self.is_amp = is_amp
        self.max_grad_norm = max_grad_norm
        self.scaler = scaler

        best_metric = 0
        bad_epochs = 0
        if snapshot_ensemble:
            test_outputs = []
            val_outputs = []

        for epoch in range(epochs):
            if bad_epochs < early_stopping_rounds:
                train_loss = self.train_fn(self.train_data_loader1, self.train_data_loader2)

                print ("Running evaluation on whole training data")
                train_out, train_loss, train_targets = self.eval_fn(self.train_data_loader1, self.train_data_loader2)
                print ("Running evaluation on evaluation data")
                val_out, val_loss, val_targets = self.eval_fn(self.val_data_loader1, self.val_data_loader2)

                train_loss = train_loss/len(self.train_data_loader1)
                val_loss = val_loss/len(self.val_data_loader1)

                #train_out = np.round(train_out)

                if snapshot_ensemble:
                    val_outputs += [np.expand_dims(np.array(val_out),axis=-1)]
                    val_out = np.concatenate(val_outputs,axis=-1).mean(axis=-1)
                #else:
                #    val_out = np.round(val_out) #np.array(outputs) >= 0.5

                train_metric = scorer(train_targets, train_out)
                val_metric = scorer(val_targets,val_out)

                with np.printoptions(precision=3):
                    print("Train loss = {} Train metric = {} Val loss = {} Val metric = {}".format(round(train_loss.detach().cpu().numpy().item(), 3),round(train_metric,3),\
                        round(val_loss.detach().cpu().numpy().item(), 3),round(val_metric,3)))

                if val_metric > best_metric:
                    #torch.save(self.model.state_dict(), os.path.join(MODEL_PATH,'model.bin'))
                    torch.save(self.model, os.path.join(MODEL_PATH,'model.bin'))
                    best_metric = val_metric
                    bad_epochs = 0

                else:
                    bad_epochs += 1

                if snapshot_ensemble and self.test_data_loader1 and self.test_data_loader2:
                    test_out = test_fn_multi(self.test_data_loader1,self.test_data_loader2, self.model, self.device, self.final_activation)

                    test_outputs += [np.expand_dims(np.array(test_out),axis=-1)]

            stats.update({"epoch_{}".format(epoch): {"train_loss": round(train_loss.detach().cpu().numpy().item(), 3), "train_metric": round(train_metric,3), \
                "val_loss": round(val_loss.detach().cpu().numpy().item(), 3), "val_metric": round(val_metric,3)}})

            if use_wandb:
                wandb.log({"train_loss": train_loss, "val_loss": val_loss, "train_metric": train_metric, "val_metric": val_metric})

        if self.test_data_loader1 and self.test_data_loader2:
            if snapshot_ensemble:
                self.test_output = np.concatenate(test_outputs,axis=-1).mean(axis=-1)
            else:
                self.test_output = test_fn_multi(self.test_data_loader1, self.test_data_loader2, self.model, self.device, self.final_activation)
        else:
            self.test_output = []

        #with np.printoptions(precision=3):
        #    print (stats)

        if use_wandb == False:
            with open(os.path.join(MODEL_PATH,'all_stats.json'), 'w') as outfile:
                json.dump(stats, outfile)

            with open(os.path.join(MODEL_PATH,'final_stats.json'), 'w') as outfile:
                d = {"epoch": epoch, "train_loss": round(train_loss.detach().cpu().numpy().item(), 3), "train_metric": round(train_metric,3), \
                "val_loss": round(val_loss.detach().cpu().numpy().item(), 3), "val_metric": round(val_metric,3)}
                json.dump(d, outfile)

class QATrainer:

    def __init__(self, model, train_data_loader, val_data_loader, device, model_description, final_activation=None, test_data_loader=None):
        self.model = model
        self.train_data_loader = train_data_loader
        self.val_data_loader = val_data_loader
        self.device = device
        self.description = model_description
        self.test_data_loader = test_data_loader
        self.final_activation = final_activation

        self.model.to(self.device)

        self.print_stats()

    def print_stats(self):
        print ("[LOG] Total number of parameters to learn {}".format(sum(p.numel() for p in self.model.parameters() \
                                                                 if p.requires_grad)))

    def train_fn(self, data_loader):
        self.model.train()

        total_loss = 0
        pbar = tqdm(enumerate(data_loader), total=len(data_loader))
        for bi, d in pbar:
            ids = d["ids"]
            token_type_ids = d["token_type_ids"]
            mask = d["mask"]
            targets_start = d["targets_start"]
            targets_end = d["targets_end"]

            ids = ids.to(self.device)
            token_type_ids = token_type_ids.to(self.device)
            mask = mask.to(self.device)
            targets_start = targets_start.to(self.device)
            targets_end = targets_end.to(self.device)

            self.optimizer.zero_grad()

            #with amp.autocast():
            outputs_start, outputs_end = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
            loss = self.loss_fn(outputs_start, outputs_end, targets_start, targets_end)

            #nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

            if self.scaler:
                #self.scaler.scale(loss.backward())
                #self.scaler.step(self.optimizer.step())
                #self.scaler.update()
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
                self.optimizer.step()
                self.scheduler.step()

            else:
                loss.backward()
                if self.num_tpus == 0:
                    self.optimizer.step()
                else:
                    xm.optimizer_step(self.optimizer, barrier=True)

                self.scheduler.step()

            total_loss += loss
            with np.printoptions(precision=3):
                v = round(loss.cpu().detach().numpy().item(), 3)
                pbar.set_description("Current training Loss {}".format(v))

        return total_loss

    def eval_fn(self, data_loader):
        self.model.eval()
        fin_targets_start = []
        fin_targets_end = []
        fin_outputs_start = []
        fin_outputs_end = []

        total_loss = 0
        with torch.no_grad():
            pbar = tqdm(enumerate(data_loader), total=len(data_loader))
            for bi, d in pbar:
                ids = d["ids"]
                token_type_ids = d["token_type_ids"]
                mask = d["mask"]
                targets_start = d["targets_start"]
                targets_end = d["targets_end"]

                ids = ids.to(self.device)
                token_type_ids = token_type_ids.to(self.device)
                mask = mask.to(self.device)
                targets_start = targets_start.to(self.device)
                targets_end = targets_end.to(self.device)

                outputs_start, outputs_end = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
                loss = self.loss_fn(outputs_start, outputs_end, targets_start, targets_end)

                total_loss += loss

                with np.printoptions(precision=3):
                    v = round(loss.cpu().detach().numpy().item(), 3)
                    pbar.set_description("Current eval Loss {}".format(v))

                fin_targets_start.extend(targets_start.cpu().detach().numpy().tolist())
                fin_targets_end.extend(targets_end.cpu().detach().numpy().tolist())

                if self.final_activation == 'sigmoid':
                    fin_outputs_start.extend(torch.sigmoid(outputs_start).cpu().detach().numpy().tolist())
                    fin_outputs_end.extend(torch.sigmoid(outputs_end).cpu().detach().numpy().tolist())
                else:
                    fin_outputs_start.extend(outputs_start.cpu().detach().numpy().tolist())
                    fin_outputs_end.extend(outputs_end.cpu().detach().numpy().tolist())

        return fin_outputs_start, fin_outputs_end, total_loss, fin_targets_start, fin_targets_end

    def train(self, epochs, lr, scorer, loss_fn, optimizer, scheduler, MODEL_PATH, num_gpus=0, num_tpus=0, max_grad_norm=1, early_stopping_rounds=3, \
                    snapshot_ensemble=False, is_amp=False, use_wandb=False, seed=42):

        self.scorer = scorer
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.num_gpus = num_gpus
        self.num_tpus = num_tpus

        set_seed(seed)

        self.args = {'description': self.description, 'loss': loss_fn.__class__.__name__, 'epochs': epochs, 'learning_rate': lr, \
                    'optimizer': self.optimizer.__class__.__name__, 'scheduler': self.scheduler.__class__.__name__}

        if use_wandb and not _torch_tpu_available:
            wandb.init(project="Project",config=self.args)
            wandb.watch(self.model)
        else:
            with open(os.path.join(MODEL_PATH, 'parameters.json'),'w') as outfile:
                json.dump(self.args, outfile)

        if _num_gpus > 1 and (num_gpus != 0 and num_gpus != 1):
            self.model = nn.DataParallel(self.model)

        if is_amp and _has_apex:
            #scaler = amp.GradScaler()
            scaler = 1
        else:
            scaler = None

        if os.path.exists(os.path.join(MODEL_PATH,'model.bin')):
            try:
                try:
                    #self.model.load_state_dict(torch.load(os.path.join(MODEL_PATH,'model.bin')))
                    self.model.load_state_dict(torch.load(os.path.join(MODEL_PATH,'model.bin')))
                except:
                    self.model = torch.load(os.path.join(MODEL_PATH,'model.bin'))

                #print (self.model)
                print ("Loaded model from previous checkpoint")
            except:
                pass

        self.is_amp = is_amp
        self.max_grad_norm = max_grad_norm
        self.scaler = scaler

        best_metric = 0
        bad_epochs = 0
        if snapshot_ensemble:
            test_outputs_start = []
            val_outputs_start = []
            test_outputs_end = []
            val_outputs_end = []

        stats = {}
        for epoch in range(epochs):
            if bad_epochs < early_stopping_rounds:
                train_loss = self.train_fn(self.train_data_loader)

                print ("Running evaluation on whole training data")
                train_out_start, train_out_end, train_loss, train_targets_start, train_targets_end = self.eval_fn(self.train_data_loader)
                print ("Running evaluation on validation data")
                val_out_start, val_out_end, val_loss, val_targets_start, val_targets_end = self.eval_fn(self.val_data_loader)

                train_loss = train_loss/len(self.train_data_loader)
                val_loss = val_loss/len(self.val_data_loader)

                #train_out = np.round(train_out)

                if snapshot_ensemble:
                    val_outputs_start += [np.expand_dims(np.array(val_out_start),axis=-1)]
                    val_outputs_end += [np.expand_dims(np.array(val_out_end),axis=-1)]

                    val_out_start = np.concatenate(val_outputs_start,axis=-1).mean(axis=-1)
                    val_out_end = np.concatenate(val_outputs_end,axis=-1).mean(axis=-1)
                #else:
                #    val_out = np.round(val_out) #np.array(outputs) >= 0.5

                train_metric = scorer(train_targets_start, train_targets_end, train_out_start, train_out_end)
                val_metric = scorer(val_targets_start, val_targets_end, val_out_start, val_out_end)

                with np.printoptions(precision=3):
                    print("Train loss = {} Train metric = {} Val loss = {} Val metric = {}".format(round(train_loss.detach().cpu().numpy().item(), 3),round(train_metric,3),\
                        round(val_loss.detach().cpu().numpy().item(), 3),round(val_metric, 3)))

                if val_metric > best_metric:
                    #torch.save(self.model.state_dict(), os.path.join(MODEL_PATH,'model.bin'))
                    torch.save(self.model, os.path.join(MODEL_PATH,'model.bin'))
                    best_metric = val_metric
                    bad_epochs = 0

                else:
                    bad_epochs += 1

                if snapshot_ensemble and self.test_data_loader:
                    test_out_start, test_out_end = test_fn_qa(self.test_data_loader, self.model, self.device, self.final_activation)

                    test_outputs_start += [np.expand_dims(np.array(test_out_start),axis=-1)]
                    test_outputs_end += [np.expand_dims(np.array(test_out_end),axis=-1)]

            stats.update({"epoch_{}".format(epoch): {"train_loss": round(train_loss.detach().cpu().numpy().item(), 3), "train_metric": round(train_metric,3), \
                "val_loss": round(val_loss.detach().cpu().numpy().item(), 3),  "val_metric": round(val_metric,3)}})

            if use_wandb:
                wandb.log({"train_loss": train_loss, "val_loss": val_loss, "train_metric": train_metric, "val_metric": val_metric})

        if self.test_data_loader:
            if snapshot_ensemble:
                self.test_output_start = np.concatenate(test_outputs_start,axis=-1).mean(axis=-1)
                self.test_output_end = np.concatenate(test_outputs_end,axis=-1).mean(axis=-1)
            else:
                self.test_output_start, self.test_output_end = test_fn_qa(self.test_data_loader, self.model, self.device, self.final_activation)
        else:
            self.test_output_start = []
            self.test_output_end = []

        #with np.printoptions(precision=3):
        #    print (stats)

        if use_wandb == False:
            with open(os.path.join(MODEL_PATH,'all_stats.json'), 'w') as outfile:
                json.dump(stats, outfile)

            with open(os.path.join(MODEL_PATH,'final_stats.json'), 'w') as outfile:
                d = {"epoch": epoch, "train_loss": round(train_loss.detach().cpu().numpy().item(), 3), "train_metric": round(train_metric,3), \
                "val_loss": round(val_loss.detach().cpu().numpy().item(), 3),  "val_metric": round(val_metric,3)}
                json.dump(d, outfile)

class PLTrainer(pl.LightningModule):

    def __init__(self, num_train_steps, model, metric, loss_fn, lr, final_activation=None, seed=42):
        super(PLTrainer, self).__init__()

        seed_everything(seed)

        self.model = model
        self.loss_fn = loss_fn #losses.get_loss(loss_fn)
        self.metric_name = metric.__class__.__name__
        self.metric = metric
        self.num_train_steps = num_train_steps
        self.lr = lr
        self.final_activation = final_activation

        self.save_hyperparameters()

        self.print_stats()

    def print_stats(self):
        print ("[LOG] Total number of parameters to learn {}".format(sum(p.numel() for p in self.model.parameters() \
                                                                 if p.requires_grad)))

    def forward(self, x):

        return self.model(ids=x["ids"], mask=x["mask"], token_type_ids=x["token_type_ids"])

    def training_step(self, batch, batch_idx):
        di = batch
        ids = di["ids"]
        token_type_ids = di["token_type_ids"]
        mask = di["mask"]
        targets = di["targets"]

        outputs = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
        if self.loss_fn.__class__.__name__ == 'masked_CELoss':
            loss = self.loss_fn(outputs, targets, mask)
        else:
            loss = self.loss_fn(outputs, targets)

        if self.final_activation == 'sigmoid':
            metric_value = self.metric(targets, torch.sigmoid(outputs))
        else:
            metric_value = self.metric(targets, outputs)

        tensorboard_logs = {'train_loss': loss, "train {}".format(self.metric_name): metric_value}

        return {'loss': loss, 'train_metric': metric_value, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        di = batch
        ids = di["ids"]
        token_type_ids = di["token_type_ids"]
        mask = di["mask"]
        targets = di["targets"]

        outputs = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
        if self.loss_fn.__class__.__name__ == 'masked_CELoss':
            loss = self.loss_fn(outputs, targets, mask)
        else:
            loss = self.loss_fn(outputs, targets)

        if self.final_activation == 'sigmoid':
            metric_value = self.metric(targets,torch.sigmoid(outputs))
        else:
            metric_value = self.metric(targets,outputs)

        tensorboard_logs = {'val_loss': loss, "val {}".format(self.metric_name): metric_value}

        return {'val_loss': loss, 'val_metric': metric_value, 'log': tensorboard_logs}

    def configure_optimizers(self):

        param_optimizer = list(self.model.named_parameters())
        no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
        optimizer_parameters = [
            {
                "params": [
                    p for n, p in param_optimizer if not any(nd in n for nd in no_decay)
                ],
                "weight_decay": 0.001,
            },
            {
                "params": [
                    p for n, p in param_optimizer if any(nd in n for nd in no_decay)
                ],
                "weight_decay": 0.0,
            },
        ]

        optimizer = AdamW(optimizer_parameters, lr=self.lr)
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=self.num_train_steps)

        return [optimizer], [{'scheduler': scheduler}]

    def training_epoch_end(self, outputs):
        train_loss_mean = torch.stack([x['loss'] for x in outputs]).mean()
        train_metric_mean = torch.stack([x['train_metric'] for x in outputs]).mean()
        print ("Train loss = {} Train metric = {}".format(round(train_loss_mean.detach().cpu().numpy().item(), 3),round(train_metric_mean.detach().cpu().numpy().item(), 3)))

        return {'train_loss': train_loss_mean, 'train_metric': train_metric_mean}

    def validation_epoch_end(self, outputs):
        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
        val_metric_mean = torch.stack([x['val_metric'] for x in outputs]).mean()
        print ("val loss = {} val metric = {} ".format(round(val_loss_mean.detach().cpu().numpy().item(), 3),round(val_metric_mean.detach().cpu().numpy().item(), 3)))

        return {'val_loss': val_loss_mean, 'val_metric': val_metric_mean}

class PLTrainerQA(pl.LightningModule):

    def __init__(self, num_train_steps, model, metric, loss_fn, lr, final_activation=None, seed=42):
        super(PLTrainerQA, self).__init__()

        seed_everything(seed)

        self.model = model
        self.loss_fn = loss_fn #losses.get_loss(loss_fn)
        self.metric_name = metric.__class__.__name__
        self.metric = metric
        self.num_train_steps = num_train_steps
        self.lr = lr
        self.final_activation = final_activation

        self.save_hyperparameters()

        self.print_stats()

    def print_stats(self):
        print ("[LOG] Total number of parameters to learn {}".format(sum(p.numel() for p in self.model.parameters() \
                                                                 if p.requires_grad)))

    def forward(self, x):

        return self.model(ids=x["ids"], mask=x["mask"], token_type_ids=x["token_type_ids"])

    def training_step(self, batch, batch_idx):
        di = batch
        ids = di["ids"]
        token_type_ids = di["token_type_ids"]
        mask = di["mask"]
        targets_start = di["targets_start"]
        targets_end = di["targets_end"]

        outputs_start, outputs_end = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
        loss = self.loss_fn(outputs_start, outputs_end, targets_start, targets_end)

        if self.final_activation == 'sigmoid':
            metric_value = self.metric(targets_start, targets_end, torch.sigmoid(outputs_start), torch.sigmoid(outputs_end))
        else:
            metric_value = self.metric(targets_start, targets_end, outputs_start, outputs_end)

        tensorboard_logs = {'train_loss': loss, "train {}".format(self.metric_name): metric_value}

        return {'loss': loss, 'train_metric': metric_value, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        di = batch
        ids = di["ids"]
        token_type_ids = di["token_type_ids"]
        mask = di["mask"]
        targets_start = di["targets_start"]
        targets_end = di["targets_end"]

        outputs_start, outputs_end = self.model(ids=ids, mask=mask, token_type_ids=token_type_ids)
        loss = self.loss_fn(outputs_start, outputs_end, targets_start, targets_end)

        if self.final_activation == 'sigmoid':
            metric_value = self.metric(targets_start, targets_end, torch.sigmoid(outputs_start), torch.sigmoid(outputs_end))
        else:
            metric_value = self.metric(targets_start, targets_end, outputs_start, outputs_end)

        tensorboard_logs = {'val_loss': loss, "val {}".format(self.metric_name): metric_value}

        return {'val_loss': loss, 'val_metric': metric_value, 'log': tensorboard_logs}

    def configure_optimizers(self):

        param_optimizer = list(self.model.named_parameters())
        no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
        optimizer_parameters = [
            {
                "params": [
                    p for n, p in param_optimizer if not any(nd in n for nd in no_decay)
                ],
                "weight_decay": 0.001,
            },
            {
                "params": [
                    p for n, p in param_optimizer if any(nd in n for nd in no_decay)
                ],
                "weight_decay": 0.0,
            },
        ]

        optimizer = AdamW(optimizer_parameters, lr=self.lr)
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=self.num_train_steps)

        return [optimizer], [{'scheduler': scheduler}]

    def training_epoch_end(self, outputs):
        train_loss_mean = torch.stack([x['loss'] for x in outputs]).mean()
        train_metric_mean = torch.stack([x['train_metric'] for x in outputs]).mean()
        print ("Train loss = {} Train metric = {}".format(round(train_loss_mean.detach().cpu().numpy().item(), 3),round(train_metric_mean.detach().cpu().numpy().item(), 3)))

        return {'train_loss': train_loss_mean, 'train_metric': train_metric_mean}

    def validation_epoch_end(self, outputs):
        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
        val_metric_mean = torch.stack([x['val_metric'] for x in outputs]).mean()
        print ("val loss = {} val metric = {} ".format(round(val_loss_mean.detach().cpu().numpy().item(), 3),round(val_metric_mean.detach().cpu().numpy().item(), 3)))

        return {'val_loss': val_loss_mean, 'val_metric': val_metric_mean}


text:
from unittest import TestCase, main
from program_graphs import CFG
from program_graphs.cfg.edge_contraction import is_possible_to_contract, edge_contraction_all
import networkx as nx  # type: ignore

class TestCFGEdgeContraction(TestCase):

    def test_is_contraction_possible_linear(self) -> None:
        cfg = CFG()
        cfg.add_edges_from([
            (1, 2), (2, 3)
        ])
        self.assertTrue(is_possible_to_contract(cfg, (1, 2)))
        self.assertTrue(is_possible_to_contract(cfg, (2, 3)))

    def test_is_contraction_possible_simple_cycle(self) -> None:
        cfg = CFG()
        cfg.add_edges_from([
            (1, 2), (2, 3), (3, 4), (3, 2)
        ])
        self.assertFalse(is_possible_to_contract(cfg, (1, 2)))
        self.assertTrue(is_possible_to_contract(cfg, (2, 3)))
        self.assertFalse(is_possible_to_contract(cfg, (3, 4)))
        self.assertFalse(is_possible_to_contract(cfg, (3, 2)))

    def test_is_contraction_possible_two_nested_cycles(self) -> None:
        cfg = CFG()
        cfg.add_edges_from([
            (1, 2), (2, 3), (3, 4), (4, 5),
            (3, 3), (4, 2)
        ])
        self.assertFalse(is_possible_to_contract(cfg, (3, 3)))
        self.assertFalse(is_possible_to_contract(cfg, (4, 2)))
        self.assertFalse(is_possible_to_contract(cfg, (1, 2)))
        self.assertFalse(is_possible_to_contract(cfg, (2, 3)))
        self.assertFalse(is_possible_to_contract(cfg, (3, 4)))
        self.assertFalse(is_possible_to_contract(cfg, (4, 5)))

    def test_is_contraction_possible_two_nested_cycles_2(self) -> None:
        cfg = CFG()
        cfg.add_edges_from([
            (1, 2), (2, 3), (3, 4), (4, 5),
            (3, 2), (4, 2)
        ])
        self.assertTrue(is_possible_to_contract(cfg, (2, 3)))

    def test_is_contraction_possible_another_cycle(self) -> None:
        cfg = CFG()
        cfg.add_edges_from([
            (1, 2), (2, 3), (3, 4), (4, 2), (2, 5)
        ])
        self.assertTrue(is_possible_to_contract(cfg, (3, 4)))
        self.assertFalse(is_possible_to_contract(cfg, (2, 3)))

    def test_is_contraction_possible_return_left(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([(node_1, node_2)])
        cfg.add_return_node(node_1)
        self.assertFalse(is_possible_to_contract(cfg, (node_1, node_2)))

    def test_is_contraction_possible_return_right(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([(node_1, node_2)])
        cfg.add_return_node(node_2)
        self.assertTrue(is_possible_to_contract(cfg, (node_1, node_2)))

    def test_is_contraction_possible_break_node_left(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([(node_1, node_2)])
        cfg.add_break_node(node_1)
        self.assertFalse(is_possible_to_contract(cfg, (node_1, node_2)))

    def test_is_contraction_possible_break_node_right(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([(node_1, node_2)])
        cfg.add_break_node(node_2)
        self.assertTrue(is_possible_to_contract(cfg, (node_1, node_2)))

    def test_is_contraction_possible_continue_node_left(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([(node_1, node_2)])
        cfg.add_continue_node(node_1)
        self.assertFalse(is_possible_to_contract(cfg, (node_1, node_2)))

    def test_is_contraction_possible_continue_node_right(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([(node_1, node_2)])
        cfg.add_continue_node(node_2)
        self.assertTrue(is_possible_to_contract(cfg, (node_1, node_2)))

    def test_edge_contraction_case_1(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        node_3 = cfg.add_node([3])
        node_4 = cfg.add_node([4])
        cfg.add_edges_from([(node_1, node_2), (node_2, node_3), (node_3, node_4)])
        cfg = edge_contraction_all(cfg)
        self.assertEqual(len(cfg.nodes()), 1)

    def test_edge_contraction_case_cycle_of_length_one(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        node_3 = cfg.add_node([3])
        node_4 = cfg.add_node([4])
        cfg.add_edges_from([(node_1, node_2), (node_2, node_3), (node_3, node_4), (node_3, node_2)])
        cfg = edge_contraction_all(cfg)
        self.assertTrue(
            nx.algorithms.is_isomorphic(cfg, nx.DiGraph([(1, 2), (2, 3), (2, 2)]))
        )

    def test_edge_contraction_case_cycle_of_length_two(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        node_3 = cfg.add_node([3])
        node_4 = cfg.add_node([4])
        node_5 = cfg.add_node([5])
        cfg.add_edges_from([(node_1, node_2), (node_2, node_3), (node_3, node_4), (node_4, node_5), (node_4, node_2)])
        cfg = edge_contraction_all(cfg)
        self.assertTrue(
            nx.algorithms.is_isomorphic(cfg, nx.DiGraph([(1, 2), (2, 3), (2, 2)]))
        )

    def test_edge_contraction_case_2(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([])
        node_2 = cfg.add_node([])
        node_3 = cfg.add_node([])
        node_4 = cfg.add_node([])
        node_5 = cfg.add_node([])
        node_6 = cfg.add_node([])
        cfg.add_edges_from([
            (node_1, node_2),
            (node_2, node_3),
            (node_3, node_4),
            (node_4, node_5),
            (node_3, node_6),
            (node_5, node_6)
        ])
        cfg = edge_contraction_all(cfg)
        self.assertTrue(
            nx.algorithms.is_isomorphic(cfg, nx.DiGraph([(1, 3), (1, 2), (2, 3)]))
        )

    def test_edge_contraction_name_assignment_left(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([], name='A')
        node_2 = cfg.add_node([])
        cfg.add_edges_from([(node_1, node_2)])
        cfg = edge_contraction_all(cfg)
        self.assertIsNotNone(cfg.find_node_by_name('A'))

    def test_edge_contraction_name_assignment_right(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([])
        node_2 = cfg.add_node([], name='B')
        cfg.add_edges_from([(node_1, node_2)])
        cfg = edge_contraction_all(cfg)
        self.assertIsNotNone(cfg.find_node_by_name('B'))

    def test_edge_contraction_name_assignment_both(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([], name='A')
        node_2 = cfg.add_node([], name='B')
        cfg.add_edges_from([(node_1, node_2)])
        cfg = edge_contraction_all(cfg)
        self.assertIsNotNone(cfg.find_node_by_name('B'))

    def test_blocks_merging(self) -> None:
        cfg = CFG()
        node_1 = cfg.add_node([1])
        node_2 = cfg.add_node([2])
        cfg.add_edges_from([
            (node_1, node_2)
        ])
        cfg = edge_contraction_all(cfg)
        self.assertIn([1, 2], cfg.node_id_2_block)

if __name__ == '__main__':
    main()

ext: py

max_stars_repo_name: acheshkov/program-graphs

max_issues_repo_name: acheshkov/program-graphs

max_forks_repo_name: acheshkov/program-graphs

reason: import networkx


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

# This code is part of Qiskit.
#
# (C) Copyright IBM 2018, 2020.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

""" Test Operator construction, including OpPrimitives and singletons. """

import unittest
from test.aqua import QiskitAquaTestCase
import itertools
import numpy as np

from qiskit import QuantumCircuit
from qiskit.quantum_info.operators import Operator, Pauli
from qiskit.circuit.library import CZGate

from qiskit.aqua.operators import X, Y, Z, I, CX, T, H, PrimitiveOp, SummedOp, PauliOp, Minus

# pylint: disable=invalid-name

class TestOpConstruction(QiskitAquaTestCase):
    """Operator Construction tests."""

    def test_pauli_primitives(self):
        """ from to file test """
        newop = X ^ Y ^ Z ^ I
        self.assertEqual(newop.primitive, Pauli(label='XYZI'))

        kpower_op = (Y ^ 5) ^ (I ^ 3)
        self.assertEqual(kpower_op.primitive, Pauli(label='YYYYYIII'))

        kpower_op2 = (Y ^ I) ^ 4
        self.assertEqual(kpower_op2.primitive, Pauli(label='YIYIYIYI'))

        # Check immutability
        self.assertEqual(X.primitive, Pauli(label='X'))
        self.assertEqual(Y.primitive, Pauli(label='Y'))
        self.assertEqual(Z.primitive, Pauli(label='Z'))
        self.assertEqual(I.primitive, Pauli(label='I'))

    def test_composed_eval(self):
        """ Test eval of ComposedOp """
        self.assertAlmostEqual(Minus.eval('1'), -.5**.5)

    def test_evals(self):
        """ evals test """
        # pylint: disable=no-member
        # TODO: Think about eval names
        self.assertEqual(Z.eval('0').eval('0'), 1)
        self.assertEqual(Z.eval('1').eval('0'), 0)
        self.assertEqual(Z.eval('0').eval('1'), 0)
        self.assertEqual(Z.eval('1').eval('1'), -1)
        self.assertEqual(X.eval('0').eval('0'), 0)
        self.assertEqual(X.eval('1').eval('0'), 1)
        self.assertEqual(X.eval('0').eval('1'), 1)
        self.assertEqual(X.eval('1').eval('1'), 0)
        self.assertEqual(Y.eval('0').eval('0'), 0)
        self.assertEqual(Y.eval('1').eval('0'), -1j)
        self.assertEqual(Y.eval('0').eval('1'), 1j)
        self.assertEqual(Y.eval('1').eval('1'), 0)

        # Check that Pauli logic eval returns same as matrix logic
        self.assertEqual(PrimitiveOp(Z.to_matrix()).eval('0').eval('0'), 1)
        self.assertEqual(PrimitiveOp(Z.to_matrix()).eval('1').eval('0'), 0)
        self.assertEqual(PrimitiveOp(Z.to_matrix()).eval('0').eval('1'), 0)
        self.assertEqual(PrimitiveOp(Z.to_matrix()).eval('1').eval('1'), -1)
        self.assertEqual(PrimitiveOp(X.to_matrix()).eval('0').eval('0'), 0)
        self.assertEqual(PrimitiveOp(X.to_matrix()).eval('1').eval('0'), 1)
        self.assertEqual(PrimitiveOp(X.to_matrix()).eval('0').eval('1'), 1)
        self.assertEqual(PrimitiveOp(X.to_matrix()).eval('1').eval('1'), 0)
        self.assertEqual(PrimitiveOp(Y.to_matrix()).eval('0').eval('0'), 0)
        self.assertEqual(PrimitiveOp(Y.to_matrix()).eval('1').eval('0'), -1j)
        self.assertEqual(PrimitiveOp(Y.to_matrix()).eval('0').eval('1'), 1j)
        self.assertEqual(PrimitiveOp(Y.to_matrix()).eval('1').eval('1'), 0)

        pauli_op = Z ^ I ^ X ^ Y
        mat_op = PrimitiveOp(pauli_op.to_matrix())
        full_basis = list(map(''.join, itertools.product('01', repeat=pauli_op.num_qubits)))
        for bstr1, bstr2 in itertools.product(full_basis, full_basis):
            # print('{} {} {} {}'.format(bstr1, bstr2, pauli_op.eval(bstr1, bstr2),
            # mat_op.eval(bstr1, bstr2)))
            np.testing.assert_array_almost_equal(pauli_op.eval(bstr1).eval(bstr2),
                                                 mat_op.eval(bstr1).eval(bstr2))

        gnarly_op = SummedOp([(H ^ I ^ Y).compose(X ^ X ^ Z).tensor(Z),
                              PrimitiveOp(Operator.from_label('+r0I')),
                              3 * (X ^ CX ^ T)], coeff=3 + .2j)
        gnarly_mat_op = PrimitiveOp(gnarly_op.to_matrix())
        full_basis = list(map(''.join, itertools.product('01', repeat=gnarly_op.num_qubits)))
        for bstr1, bstr2 in itertools.product(full_basis, full_basis):
            np.testing.assert_array_almost_equal(gnarly_op.eval(bstr1).eval(bstr2),
                                                 gnarly_mat_op.eval(bstr1).eval(bstr2))

    def test_circuit_construction(self):
        """ circuit construction test """
        hadq2 = H ^ I
        cz = hadq2.compose(CX).compose(hadq2)
        qc = QuantumCircuit(2)
        qc.append(cz.primitive, qargs=range(2))

        ref_cz_mat = PrimitiveOp(CZGate()).to_matrix()
        np.testing.assert_array_almost_equal(cz.to_matrix(), ref_cz_mat)

    def test_io_consistency(self):
        """ consistency test """
        new_op = X ^ Y ^ I
        label = 'XYI'
        # label = new_op.primitive.to_label()
        self.assertEqual(str(new_op.primitive), label)
        np.testing.assert_array_almost_equal(new_op.primitive.to_matrix(),
                                             Operator.from_label(label).data)
        self.assertEqual(new_op.primitive, Pauli(label=label))

        x_mat = X.primitive.to_matrix()
        y_mat = Y.primitive.to_matrix()
        i_mat = np.eye(2, 2)
        np.testing.assert_array_almost_equal(new_op.primitive.to_matrix(),
                                             np.kron(np.kron(x_mat, y_mat), i_mat))

        hi = np.kron(H.to_matrix(), I.to_matrix())
        hi2 = Operator.from_label('HI').data
        hi3 = (H ^ I).to_matrix()
        np.testing.assert_array_almost_equal(hi, hi2)
        np.testing.assert_array_almost_equal(hi2, hi3)

        xy = np.kron(X.to_matrix(), Y.to_matrix())
        xy2 = Operator.from_label('XY').data
        xy3 = (X ^ Y).to_matrix()
        np.testing.assert_array_almost_equal(xy, xy2)
        np.testing.assert_array_almost_equal(xy2, xy3)

        # Check if numpy array instantiation is the same as from Operator
        matrix_op = Operator.from_label('+r')
        np.testing.assert_array_almost_equal(PrimitiveOp(matrix_op).to_matrix(),
                                             PrimitiveOp(matrix_op.data).to_matrix())
        # Ditto list of lists
        np.testing.assert_array_almost_equal(PrimitiveOp(matrix_op.data.tolist()).to_matrix(),
                                             PrimitiveOp(matrix_op.data).to_matrix())

        # TODO make sure this works once we resolve endianness mayhem
        # qc = QuantumCircuit(3)
        # qc.x(2)
        # qc.y(1)
        # from qiskit import BasicAer, QuantumCircuit, execute
        # unitary = execute(qc, BasicAer.get_backend('unitary_simulator')).result().get_unitary()
        # np.testing.assert_array_almost_equal(new_op.primitive.to_matrix(), unitary)

    def test_to_matrix(self):
        """to matrix text """
        np.testing.assert_array_equal(X.to_matrix(), Operator.from_label('X').data)
        np.testing.assert_array_equal(Y.to_matrix(), Operator.from_label('Y').data)
        np.testing.assert_array_equal(Z.to_matrix(), Operator.from_label('Z').data)

        op1 = Y + H
        np.testing.assert_array_almost_equal(op1.to_matrix(), Y.to_matrix() + H.to_matrix())

        op2 = op1 * .5
        np.testing.assert_array_almost_equal(op2.to_matrix(), op1.to_matrix() * .5)

        op3 = (4 - .6j) * op2
        np.testing.assert_array_almost_equal(op3.to_matrix(), op2.to_matrix() * (4 - .6j))

        op4 = op3.tensor(X)
        np.testing.assert_array_almost_equal(op4.to_matrix(),
                                             np.kron(op3.to_matrix(), X.to_matrix()))

        op5 = op4.compose(H ^ I)
        np.testing.assert_array_almost_equal(op5.to_matrix(), np.dot(op4.to_matrix(),
                                                                     (H ^ I).to_matrix()))

        op6 = op5 + PrimitiveOp(Operator.from_label('+r').data)
        np.testing.assert_array_almost_equal(
            op6.to_matrix(), op5.to_matrix() + Operator.from_label('+r').data)

    def test_adjoint(self):
        """ adjoint test """
        gnarly_op = 3 * (H ^ I ^ Y).compose(X ^ X ^ Z).tensor(T ^ Z) + \
            PrimitiveOp(Operator.from_label('+r0IX').data)
        np.testing.assert_array_almost_equal(np.conj(np.transpose(gnarly_op.to_matrix())),
                                             gnarly_op.adjoint().to_matrix())

    def test_primitive_strings(self):
        """ get primitives test """
        self.assertEqual(X.primitive_strings(), {'Pauli'})

        gnarly_op = 3 * (H ^ I ^ Y).compose(X ^ X ^ Z).tensor(T ^ Z) + \
            PrimitiveOp(Operator.from_label('+r0IX').data)
        self.assertEqual(gnarly_op.primitive_strings(), {'QuantumCircuit', 'Matrix'})

    def test_to_pauli_op(self):
        """ Test to_pauli_op method """
        gnarly_op = 3 * (H ^ I ^ Y).compose(X ^ X ^ Z).tensor(T ^ Z) + \
            PrimitiveOp(Operator.from_label('+r0IX').data)
        mat_op = gnarly_op.to_matrix_op()
        pauli_op = gnarly_op.to_pauli_op()
        self.assertIsInstance(pauli_op, SummedOp)
        for p in pauli_op:
            self.assertIsInstance(p, PauliOp)
        np.testing.assert_array_almost_equal(mat_op.to_matrix(), pauli_op.to_matrix())

if __name__ == '__main__':
    unittest.main()

reason: import numpy


text:
'''
  A simple Conv3D example with Keras
'''
import keras
from keras.models import Sequential
from keras.layers import Dense, Flatten, Conv3D, MaxPooling3D, Dropout
from keras.utils import to_categorical
import h5py
import numpy as np
import matplotlib.pyplot as plt

# -- Preparatory code --
# Model configuration
batch_size = 100
no_epochs = 30
learning_rate = 0.001
no_classes = 10
validation_split = 0.2
verbosity = 1

# Convert 1D vector into 3D values, provided by the 3D MNIST authors at
# https://www.kaggle.com/daavoo/3d-mnist
def array_to_color(array, cmap="Oranges"):
  s_m = plt.cm.ScalarMappable(cmap=cmap)
  return s_m.to_rgba(array)[:,:-1]

# Reshape data into format that can be handled by Conv3D layers.
# Courtesy of Sam Berglin; Zheming Lian; Jiahui Jang - University of Wisconsin-Madison
# Report - https://github.com/sberglin/Projects-and-Papers/blob/master/3D%20CNN/Report.pdf
# Code - https://github.com/sberglin/Projects-and-Papers/blob/master/3D%20CNN/network_final_version.ipynb
def rgb_data_transform(data):
  data_t = []
  for i in range(data.shape[0]):
    data_t.append(array_to_color(data[i]).reshape(16, 16, 16, 3))
  return np.asarray(data_t, dtype=np.float32)

# -- Process code --
# Load the HDF5 data file
with h5py.File("./full_dataset_vectors.h5", "r") as hf:

    # Split the data into training/test features/targets
    X_train = hf["X_train"][:]
    targets_train = hf["y_train"][:]
    X_test = hf["X_test"][:]
    targets_test = hf["y_test"][:]

    # Determine sample shape
    sample_shape = (16, 16, 16, 3)

    # Reshape data into 3D format
    X_train = rgb_data_transform(X_train)
    X_test = rgb_data_transform(X_test)

    # Convert target vectors to categorical targets
    targets_train = to_categorical(targets_train).astype(np.integer)
    targets_test = to_categorical(targets_test).astype(np.integer)

    # Create the model
    model = Sequential()
    model.add(Conv3D(32, kernel_size=(3, 3, 3), activation='relu', kernel_initializer='he_uniform', input_shape=sample_shape))
    model.add(MaxPooling3D(pool_size=(2, 2, 2)))
    model.add(Dropout(0.5))
    model.add(Conv3D(64, kernel_size=(3, 3, 3), activation='relu', kernel_initializer='he_uniform'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2)))
    model.add(Dropout(0.5))
    model.add(Flatten())
    model.add(Dense(256, activation='relu', kernel_initializer='he_uniform'))
    model.add(Dense(no_classes, activation='softmax'))

    # Compile the model
    model.compile(loss=keras.losses.categorical_crossentropy,
                  optimizer=keras.optimizers.Adam(lr=learning_rate),
                  metrics=['accuracy'])

    # Fit data to model
    history = model.fit(X_train, targets_train,
                batch_size=batch_size,
                epochs=no_epochs,
                verbose=verbosity,
                validation_split=validation_split)

    # Generate generalization metrics
    score = model.evaluate(X_test, targets_test, verbose=0)
    print(f'Test loss: {score[0]} / Test accuracy: {score[1]}')

    # Plot history: Categorical crossentropy & Accuracy
    plt.plot(history.history['loss'], label='Categorical crossentropy (training data)')
    plt.plot(history.history['val_loss'], label='Categorical crossentropy (validation data)')
    plt.plot(history.history['accuracy'], label='Accuracy (training data)')
    plt.plot(history.history['val_accuracy'], label='Accuracy (validation data)')
    plt.title('Model performance for 3D MNIST Keras Conv3D example')
    plt.ylabel('Loss value')
    plt.xlabel('No. epoch')
    plt.legend(loc="upper left")
    plt.show()

max_stars_repo_name: soufiomario/keras-cnn

max_issues_repo_name: soufiomario/keras-cnn

max_forks_repo_name: soufiomario/keras-cnn


text:
from pypot.creatures import PoppyHumanoid as PH
from pypot.sensor import OpenCVCamera
import matplotlib.pyplot as plt
import time
import numpy as np
import sys

jlims = {
    "head_y": [-22., 20.], # backward/forward
    "head_z": [-58., 40.], # right/left
    "l_ankle_y": [-40, 40], # flex/point
    "r_ankle_y": [-40, 40], # flex/point
}
fps = 10
# moving_speed = 25

p = PH()
for m in jlims.keys():
    getattr(p,m).compliant=False
    # getattr(p,m).moving_speed = moving_speed

p.goto_position({"r_ankle_y":-40,"l_ankle_y":40}, 1, wait=True)

c = OpenCVCamera("poppy-cam",0,fps)

def handle_close(event):
    c.close()
    p.close()
    sys.exit(0)

plt.ion()
fig = plt.figure()
fig.canvas.mpl_connect("close_event", handle_close)

while True:

    if p.r_ankle_y.present_position > +35 or p.l_ankle_y.present_position < -35:
        p.r_ankle_y.goto_position(-40, .8, wait=False)
        p.l_ankle_y.goto_position(+40, .8, wait=False)
    if p.r_ankle_y.present_position < -35 or p.l_ankle_y.present_position > +35:
        p.r_ankle_y.goto_position(+40, .8, wait=False)
        p.l_ankle_y.goto_position(-40, .8, wait=False)

    plt.pause(.5)
    frame = c.frame.astype(float)/255.
    plt.pause(2./fps)

    new_frame = c.frame.astype(float)/255.
    diff = np.fabs(new_frame - frame)
    frame = new_frame

    i = np.arange(diff.shape[0])[:,np.newaxis,np.newaxis]
    j = np.arange(diff.shape[1])[np.newaxis,:,np.newaxis]
    ds = diff.sum()
    ci = (diff*i).sum() / ds
    cj = (diff*j).sum() / ds

    plt.clf()
    plt.subplot(1,2,1)
    plt.imshow(frame[:,:,[2,1,0]])
    plt.subplot(1,2,2)
    plt.imshow(diff[:,:,[2,1,0]])
    plt.scatter([cj],[ci],c='r')
    plt.draw()

    ### constant velocity based on sign
    di = ci - diff.shape[0]/2
    dj = cj - diff.shape[1]/2

    dirs = {"head_y": +np.sign(di), "head_z": -np.sign(dj)}

    targets = {m: min(max(
        getattr(p,m).present_position + 15*dirs[m],
        jlims[m][0]), jlims[m][1])
        for m in dirs.keys()}

    ### absolute position based on centroid
    # targets = {
    #     "head_y": jlims["head_y"][0] + ci/diff.shape[0]*(jlims["head_y"][1] - jlims["head_y"][0]),
    #     "head_z": jlims["head_z"][1] + cj/diff.shape[1]*(jlims["head_z"][0] - jlims["head_z"][1]),
    # }

    ### velocity based on centroid
    # di = (ci - diff.shape[0]/2) / (diff.shape[0]/2)
    # dj = (cj - diff.shape[1]/2) / (diff.shape[1]/2)

    # dirs = {"head_y": +di, "head_z": -dj}

    # targets = {m: min(max(
    #     getattr(p,m).present_position + 50*dirs[m],
    #     jlims[m][0]), jlims[m][1])
    #     for m in dirs.keys()}

    p.goto_position(targets, 2, wait=True)

c.close()
p.close()

max_stars_repo_name: garrettkatz/poppy-simulations

max_issues_repo_name: garrettkatz/poppy-simulations

max_forks_repo_name: garrettkatz/poppy-simulations


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 17 18:14:37 2021
Edited Thu Nov 4 2021: rename and make this the main resource file for the rpo

@author: lizz
"""
import numpy as np
from scipy import optimize

def segments_fit(X, Y, maxcount):
    xmin = X.min()
    xmax = X.max()

    n = len(X)

    AIC_ = float('inf')
    BIC_ = float('inf')
    r_   = None

    for count in range(1, maxcount+1):

        seg = np.full(count - 1, (xmax - xmin) / count)

        px_init = np.r_[np.r_[xmin, seg].cumsum(), xmax]
        py_init = np.array([Y[np.abs(X - x) < (xmax - xmin) * 0.1].mean() for x in px_init])

        def func(p):
            seg = p[:count - 1]
            py = p[count - 1:]
            px = np.r_[np.r_[xmin, seg].cumsum(), xmax]
            return px, py

        def err(p): # This is RSS / n
            px, py = func(p)
            Y2 = np.interp(X, px, py)
            return np.mean((Y - Y2)**2)

        r = optimize.minimize(err, x0=np.r_[seg, py_init], method='Nelder-Mead')

        # Compute AIC/ BIC.
        AIC = n * np.log10(err(r.x)) + 4 * count
        BIC = n * np.log10(err(r.x)) + 2 * count * np.log(n)

        if((BIC < BIC_) & (AIC < AIC_)): # Continue adding complexity.
            r_ = r
            AIC_ = AIC
            BIC_ = BIC
        else: # Stop.
            count = count - 1
            break

    return func(r_.x) ## Return the last (n-1)

max_stars_repo_name: ehultee/stoch-SMB

max_issues_repo_name: ehultee/stoch-SMB

max_forks_repo_name: ehultee/stoch-SMB

reason: import numpy,from scipy


text:
import numpy
import random

def clip(x, lo, hi):
	if x < lo:
		return lo
	elif x > hi:
		return hi
	else:
		return x

class Crop(object):
	def __init__(self, params, data_types):
		# width and height are fractions between 0 and 1
		# they indicate the relative width/height of the crop to the original dimensions
		def parse_fraction(s):
			if '/' in s:
				parts = s.split('/')
				return (int(parts[0]), int(parts[1]))
			else:
				return float(s), 1
		self.Width = parse_fraction(params['Width'])
		self.Height = parse_fraction(params['Height'])
		self.data_types = data_types

		# this augmentation should be applied in the Dataset
		self.pre_torch = True

	def forward(self, batch):
		# pick offsets to crop
		# we use same offsets for entire batch
		xoff = random.random() * (1 - self.Width[0]/self.Width[1])
		yoff = random.random() * (1 - self.Height[0]/self.Height[1])

		# transform coordinates in image to coordinates in cropped image
		def coord_transform(x, y):
			x -= xoff
			y -= yoff
			x *= self.Width[1]/self.Width[0]
			y *= self.Height[1]/self.Height[0]
			return x, y

		for i, inputs in enumerate(batch):
			if self.data_types[i] == 'image':
				width, height = inputs[0].shape[2], inputs[0].shape[1]
				target_width = int(width * self.Width[0]) // self.Width[1]
				target_height = int(height * self.Height[0]) // self.Height[1]
				sx = int(xoff * width)
				sy = int(yoff * height)
				ex = sx + target_width
				ey = sy + target_height

				ninputs = []
				for input in inputs:
					input = input[:, sy:ey, sx:ex]
					ninputs.append(input)
				batch[i] = ninputs

			elif self.data_types[i] == 'shape':
				# shapes can be arbitrary polygons
				# so we don't really want to support it
				raise Exception('crop is not supported on shape types')

			elif self.data_types[i] == 'detection':
				# we clip detection bboxes inside the crop rectangle
				# but we also remove ones that are entirely outside the crop rectangle
				ninputs = []
				for input in inputs:
					ndetections = []
					for d in input['detections']:
						cls, sx, sy, ex, ey = d
						sx, sy = coord_transform(sx, sy)
						ex, ey = coord_transform(ex, ey)
						if ex < 0 or sx >= 1 or ey < 0 or sy >= 1:
							continue
						sx = clip(sx, 0, 1)
						sy = clip(sy, 0, 1)
						ex = clip(ex, 0, 1)
						ey = clip(ey, 0, 1)
						ndetections.append([cls, sx, sy, ex, ey])
					ndetections = numpy.array(ndetections, dtype='float32')
					ninputs.append({
						'counts': ndetections.shape[0],
						'detections': ndetections,
					})
				batch[i] = ninputs

			# we assume that other data types like int, floats, etc. don't need to be cropped

		return batch


text:
"""
Lists functions checked by test_check_functions, and functions
that are in gsw_check_functions but are not in gsw.
"""

import os

import numpy as np

import gsw
from check_functions import parse_check_functions

root_path = os.path.abspath(os.path.dirname(__file__))

# Function checks that we can't handle automatically yet.
blacklist = ['deltaSA_atlas',  # the test is complicated; doesn't fit the pattern.
             ]

d = dir(gsw)
funcnames = [name for name in d if '__' not in name]

mfuncs_all = parse_check_functions(os.path.join(root_path,
                                        'gsw_check_functions_save.m'))
mfuncs = [mf for mf in mfuncs_all if mf.name in d and mf.name not in blacklist]
mfuncnames = sorted([mf.name for mf in mfuncs])

missingnames = [mf for mf in mfuncs_all if mf.name not in d]
missingnames = sorted([mf.name for mf in missingnames])

print('Functions being checked:')
for i, name in enumerate(mfuncnames):
    print(i, name)

print('Functions not in gsw:')
for i, name in enumerate(missingnames):
    print(i, name)

max_stars_repo_name: ocesaulo/GSW-Python

max_issues_repo_name: ocesaulo/GSW-Python

max_forks_repo_name: ocesaulo/GSW-Python


text:
import tensorflow_hub as hub
import tensorflow as tf
import numpy as np
from typing import Dict, List
from tqdm import tqdm

if tf.__version__[0] == '1':
    tf.compat.v1.enable_eager_execution()

class Generator:
    """Generate text embedding
    """
    def __init__(self, model_url: str):
        """Generate text embeddings using any desired tensorflow_hub module

        Parameters
        ----------
        model_url : str
            Path of tensorflow_hub module to be used. It can be either a locally downloaded path or a direct url
        """
        self.model_url = model_url
        # self.model_url = config['embedding']['model_url']
        if 'https' in model_url:
            print("[WARNING]...You are using URL to load embedding generator backend & will take more time to load.\nThe module will be downloaded first (if not already) & then loaded. This is even true after every reboot of system. ")
        print("[INFO]...Loading embedding generator backend")
        self.model = hub.load(model_url)
        print (f"[INFO]...Module {model_url} loaded")

    def unit_generator(self, input):
        """Generate text embedding for single or unit example

        Parameters
        ----------
        input : str
            A single doc to extract text embedding
        """
        return self.model(input)

    def batch_generator(self, docs: List):
        """Generate text embeddings for batches of docs
        Note: Here a doc can be a word, sentence or even a small paragraph

        Parameters
        ----------
        docs : List
            list of all docs or sentences

        Returns
        -------
        numpy.ndarray
            numpy array with shape of (n, 512), where n = len(docs)
        """
        print("[INFO]...Extracting embeddings on batch of {len(docs)}")
        embd_array = np.array([])
        embd_array = embd_array.reshape(-1, 512)
        for doc in tqdm(docs):
            _unit_embd = Generator.unit_generator([doc])
            _unit_embd_array = np.asarray(_unit_embd)
            embd_array = np.append(embd_array,_unit_embd_array, axis=0)
        return embd_array


text:
# Copyright (c) 2022, Skolkovo Institute of Science and Technology (Skoltech)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import pytest

import evops.metrics.constants
from evops.utils.IoUOverlap import __is_overlapped_iou

def test_full_iou_overlap():
    evops.metrics.constants.IOU_THRESHOLD = 0.75

    pred_indices = np.array([1, 2, 3, 4])
    gt_indices = np.array([1, 2, 3, 4])

    assert __is_overlapped_iou(pred_indices, gt_indices)

def test_null_iou_overlap():
    evops.metrics.constants.IOU_THRESHOLD = 0.75

    pred_indices = np.array([1, 2, 3, 4])
    gt_indices = np.array([5, 6, 7, 8])

    assert not __is_overlapped_iou(pred_indices, gt_indices)

def test_iou_overlap():
    evops.metrics.constants.IOU_THRESHOLD = 0.25

    pred_indices = np.array([1, 2, 3, 4])
    gt_indices = np.array([1, 2, 5, 6])

    assert __is_overlapped_iou(pred_indices, gt_indices)

max_stars_repo_name: pmokeev/plane-benchmark

max_issues_repo_name: pmokeev/plane-benchmark

max_forks_repo_name: pmokeev/plane-benchmark


text:
#IA created to be integrated as a player in the PyRat Game.
#Winner of the competition between AI's in the first semester of 2019.
#This AI combines a Supervised Approach with a CGT approach in order to obtain maximum performance,
#not only against the greedy algorithm but also against other types of AI.
# The supervised learning model was trained using games of the greedy algorithm against itself,
#but also against a reinforcement learning algorithm trained against the greedy, in order to increase robustness.
#Please cite as you use this code.

import numpy as np
import random as rd
import pickle
import time
import joblib

MOVE_DOWN = 'D'
MOVE_LEFT = 'L'
MOVE_RIGHT = 'R'
MOVE_UP = 'U'
MIN_CHEESE_GT = 12

global model

# new location after a move
def move(location, move):
    if move == MOVE_UP:
        return (location[0], location[1] + 1)
    if move == MOVE_DOWN:
        return (location[0], location[1] - 1)
    if move == MOVE_LEFT:
        return (location[0] - 1, location[1])
    if move == MOVE_RIGHT:
        return (location[0] + 1, location[1])

#opponnent stategy
def distance(la, lb):
    ax,ay = la
    bx,by = lb
    return abs(bx - ax) + abs(by - ay)

#simulates opponent
def turn_of_opponent(opponentLocation, piecesOfCheese):
    closest_poc = (-1,-1)
    best_distance = -1
    for poc in piecesOfCheese:
        if distance(poc, opponentLocation) < best_distance or best_distance == -1:
            best_distance = distance(poc, opponentLocation)
            closest_poc = poc
    ax, ay = opponentLocation
    bx, by = closest_poc
    if bx > ax:
        return MOVE_RIGHT
    if bx < ax:
        return MOVE_LEFT
    if by > ay:
        return MOVE_UP
    if by < ay:
        return MOVE_DOWN
    pass

# With this template, we integrate two ML approaches to acquire robustness against the greedy algorithm
#but also against other AI's

TEAM_NAME = "GoGoRat"

# We do not need preprocessing, so we let this function empty
def preprocessing(mazeMap, mazeWidth, mazeHeight, playerLocation, opponentLocation, piecesOfCheese, timeAllowed):
    global model
    ### Here we load the previously trained model
    model = joblib.load('trained_classifier_go.pkl')

# We use a recursive function that goes through the trees of possible plays
# It takes as arguments a given situation, and return a best target piece of cheese for the player, such that aiming to grab this piece of cheese will eventually lead to a maximum score. It also returns the corresponding score
def best_target(playerLocation, opponentLocation, playerScore, opponentScore, piecesOfCheese):

    # First we should check how many pieces of cheese each player has to see if the play is over. It is the case if no pieces of cheese are left, or if playerScore or opponentScore is more than half the total number playerScore + opponentScore + piecesOfCheese
    totalPieces = len(piecesOfCheese) + playerScore + opponentScore
    if playerScore > totalPieces / 2 or opponentScore > totalPieces / 2 or len(piecesOfCheese) == 0:
        return (-1,-1), playerScore

    # If the play is not over, then the player can aim for any of the remaining pieces of cheese
    # So we will simulate the game to each of the pieces, which will then by recurrence test all
    # the possible trees.

    best_score_so_far = -1
    best_target_so_far = (-1,-1)
    for target in piecesOfCheese:
        end_state = simulate_game_until_target(
            target,playerLocation,opponentLocation,
            playerScore,opponentScore,piecesOfCheese.copy())
        _, score = best_target(*end_state)
        if score > best_score_so_far:
            best_score_so_far = score
            best_target_so_far = target

    return best_target_so_far, best_score_so_far

# Move the agent on the labyrinth using function move from aux and the different directions
# It suffices to move in the direction of the target.
# You should only run function move once and you can't move diagonally.
## Without loss of generality, we can suppose it gets there moving vertically first then horizontally

def updatePlayerLocation(target,playerLocation):
    if playerLocation[1] != target[1]:
        if target[1] < playerLocation[1]:
            playerLocation = move(playerLocation, MOVE_DOWN)
        else:
            playerLocation = move(playerLocation, MOVE_UP)
    elif target[0] < playerLocation[0]:
        playerLocation = move(playerLocation, MOVE_LEFT)
    else:
        playerLocation = move(playerLocation, MOVE_RIGHT)
    return playerLocation

#CHECK IF EITHER/BOTH PLAYERS ARE ON THE SAME SQUARE OF A CHEESE.
#If that is the case you have to remove the cheese from the piecesOfCheese list and
#add points to the score. The players get 1 point if they are alone on the square with a cheese.
#If both players are in the same square and there is a cheese on the square each player gets 0.5 points.
def checkEatCheese(playerLocation,opponentLocation,playerScore,opponentScore,piecesOfCheese):
    if playerLocation in piecesOfCheese and playerLocation == opponentLocation:
        playerScore = playerScore + 0.5
        opponentScore = opponentScore + 0.5
        piecesOfCheese.remove(playerLocation)
    else:
        if playerLocation in piecesOfCheese:
            playerScore = playerScore + 1
            piecesOfCheese.remove(playerLocation)
        if opponentLocation in piecesOfCheese:
            opponentScore = opponentScore + 1
            piecesOfCheese.remove(opponentLocation)
    return playerScore,opponentScore

#In this function we simulate what will happen until we reach the target
#You should use the two functions defined before
def simulate_game_until_target(target,playerLocation,opponentLocation,playerScore,opponentScore,piecesOfCheese):

    #While the target cheese has not yet been eaten by either player
    #We simulate how the game will evolve until that happens
    while target in piecesOfCheese:
        #Update playerLocation (position of your player) using updatePlayerLocation
        playerLocation = updatePlayerLocation(target,playerLocation)
        #Every time that we move the opponent also moves. update the position of the opponent using turn_of_opponent and move
        opponentLocation = move(opponentLocation, turn_of_opponent(opponentLocation, piecesOfCheese))
        #Finally use the function checkEatCheese to see if any of the players is in the same square of a cheese.
        playerScore, opponentScore = checkEatCheese(
            playerLocation,opponentLocation,playerScore,opponentScore,piecesOfCheese)
    return playerLocation,opponentLocation,playerScore,opponentScore,piecesOfCheese

# During our turn we continue going to the next target, unless the piece of cheese it originally contained has been taken
# In such case, we compute the new best target to go to
current_target = (-1,-1)
def turn(mazeMap, mazeWidth, mazeHeight, playerLocation, opponentLocation, playerScore, opponentScore, piecesOfCheese, timeAllowed):
    global current_target
    if(len(piecesOfCheese)<MIN_CHEESE_GT):
        if current_target not in piecesOfCheese:
            current_target, score = best_target(playerLocation, opponentLocation, playerScore, opponentScore, piecesOfCheese)

        if current_target[1] > playerLocation[1]:
            return MOVE_UP
        if current_target[1] < playerLocation[1]:
            return MOVE_DOWN
        if current_target[0] > playerLocation[0]:
            return MOVE_RIGHT
        return MOVE_LEFT
    else:
        return turn_supervised(mazeMap, mazeWidth, mazeHeight, playerLocation, opponentLocation, playerScore, opponentScore, piecesOfCheese, timeAllowed)

def convert_input_2(player, maze, opponent, mazeHeight, mazeWidth, piecesOfCheese):
	# We will consider twice the size of the maze to simplify the creation of the canvas
	# The canvas is initialized as a numpy tensor with 3 modes (meaning it is indexed using three integers), the third one corresponding to "layers" of the canvas.
	# Here, we just use one layer, but you can defined other ones to put more information on the play (e.g. the location of the opponent could be put in a second layer)

    im_size = (2*mazeHeight-1,2*mazeWidth-1,1)

    # We initialize a canvas with only zeros
    canvas = np.zeros(im_size)

    (x,y) = player

    # fill in the first layer of the canvas with the value 1 at the location of the cheeses, relative to the position of the player (i.e. the canvas is centered on the player location)
    center_x, center_y = mazeWidth-1, mazeHeight-1
    for (x_cheese,y_cheese) in piecesOfCheese:
        canvas[y_cheese+center_y-y,x_cheese+center_x-x,0] = 1
    return canvas

def turn_supervised(mazeMap, mazeWidth, mazeHeight, playerLocation, opponentLocation, playerScore, opponentScore, piecesOfCheese, timeAllowed):
    global model,input_tm1, action, score

    # Transform the input into the canvas using convert_input

    input_t = convert_input_2(playerLocation, mazeMap, opponentLocation, mazeHeight, mazeWidth, piecesOfCheese)

    # Predict the next action using the trained model
    output = model.predict(input_t.reshape(1,-1))
    action = output[0]
    # Return the action to perform
    return [MOVE_LEFT, MOVE_RIGHT, MOVE_UP, MOVE_DOWN][action]

def postprocessing (mazeMap, mazeWidth, mazeHeight, playerLocation, opponentLocation, playerScore, opponentScore, piecesOfCheese, timeAllowed):
    pass

max_stars_repo_name: talesmarra/GoGoRat

max_issues_repo_name: talesmarra/GoGoRat

max_forks_repo_name: talesmarra/GoGoRat


text:
#! /usr/bin/env python3
#DESCRIPTION: Get final predictions for the best CNN model
#DEPENDS: data/intermediate/sentences_ml.csv, data/tmp/w2v_320d
#CREATES: data/intermediate/cnn_predictions.csv

import csv
import logging
import random
from pathlib import Path

import numpy as np
import pandas as pd
import deeplib as lib
from keras import backend as keras_backend
import os

import tensorflow as tf
import os

lib.cudnn_error_workaround()

# best settings according to grid search:
settings = dict(
    train_embedding=True,
    n_hidden=1,
    depth_hidden=64,
    learning_rate=.004,
    epochs=4,
    loss='mean_absolute_error',
    output_dim=2,
    batch_size=128,
    )
N_REP = 10

logging.basicConfig(level=logging.INFO, format='[%(asctime)s %(name)-12s %(levelname)-5s] %(message)s')

data_root = Path.cwd()/"data"
data_file = data_root/"intermediate"/"sentences_ml.csv"
embeddings_file = data_root/"tmp"/"w2v_320d"
output_file = data_root/"intermediate"/"cnn_predictions.csv"

logging.info("Loading data")

data = pd.read_csv(data_file)
# drop gold, randomize rows, select train/val

labels = lib.encode_labels(data.value, output_dim=settings['output_dim'])

# Tokenize on test data as well because the word_index is also used to initialize the embeddings
# (and since actual words are never used this is permitted, alternative would be to rewrite the tokenization
#  to use vocabulary from embeddings instead, but has the same result for more work)
tokens, word_index = lib.tokenize(data.lemmata)

train_data = tokens[data.gold == 0]
test_data = tokens[data.gold == 1]

train_labels = labels[data.gold == 0]
test_labels = labels[data.gold == 1]

test_ids = data.id[data.gold == 1].values

logging.info("Loading embeddings")
embeddings = lib.embeddings_matrix(word_index, str(embeddings_file))

logging.info(f"Training model on training data {train_data.shape}")

# Note: This doesn't seem to actually make it deterministic :-(
# Code from https://github.com/NVIDIA/tensorflow-determinism
os.environ['TF_DETERMINISTIC_OPS'] = '1'
random_seed = 1
np.random.seed(random_seed)
random.seed(random_seed)
tf.random.set_seed(random_seed)
config = tf.compat.v1.ConfigProto(intra_op_parallelism_threads = 1,
                                  inter_op_parallelism_threads = 1)
tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))

model = lib.cnn_model(settings=settings,
                      max_sequence_length=train_data.shape[1],
                      embeddings_matrix=embeddings)

accs = []
with output_file.open('w') as outf:
    w = csv.writer(outf)
    w.writerow(["id", "repetition", "value", "confidence"])

    for it in range(N_REP):

        model.fit(train_data, train_labels, epochs=settings['epochs'], batch_size=settings['batch_size'], verbose=0)

        output = model.predict(test_data)
        keras_backend.clear_session()

        p = list(lib.predict(output))
        conf = list(lib.confidence(output))
        # Also 'predict' the true labels to convert from N neurons to single value
        actual = lib.predict(test_labels)
        acc = sum([x == y for (x, y) in zip(p, actual)]) / len(p)
        accs.append(acc)
        logging.info(f"Iteration {it}/{N_REP}, accuracy: {acc}")

        for id, pred, c in zip(test_ids, p, conf):
            w.writerow([id, it, pred, c])

logging.info(f"Done, overall accuracy {sum(accs)/len(accs)}")

max_stars_repo_name: vanatteveldt/ecosent

max_issues_repo_name: vanatteveldt/ecosent

max_forks_repo_name: vanatteveldt/ecosent