text:
# Utilities supporting gaze calibration

import numpy as np
import cv2
import pandas as pd
import os

def onoff_from_binary(data, return_duration=True):
    """Converts a binary variable data into onsets, offsets, and optionally durations

    This may yield unexpected behavior if the first value of `data` is true.

    Parameters
    ----------
    data : array-like, 1D
        binary array from which onsets and offsets should be extracted

    """
    data = data.astype(np.float).copy()
    ddata = np.hstack([[0], np.diff(data)])
    (onsets,) = np.nonzero(ddata > 0)
    # print(onsets)
    (offsets,) = np.nonzero(ddata < 0)
    # print(offsets)
    onset_first = onsets[0] < offsets[0]
    len(onsets) == len(offsets)

    on_at_end = False
    on_at_start = False
    if onset_first:
        if len(onsets) > len(offsets):
            offsets = np.hstack([offsets, [-1]])
            on_at_end = True
    else:
        if len(offsets) > len(onsets):
            onsets = np.hstack([-1, offsets])
            on_at_start = True
    onoff = np.vstack([onsets, offsets])
    if return_duration:
        duration = offsets - onsets
        if on_at_end:
            duration[-1] = len(data) - onsets[-1]
        if on_at_start:
            duration[0] = offsets[0] - 0
        onoff = np.vstack([onoff, duration])

    onoff = onoff.T.astype(np.int)
    return onoff

def time_to_index(onsets_offsets, timeline):
    """find indices between onsets & offsets in timeline

    Parameters
    ----------
    """
    out = np.zeros_like(onsets_offsets)
    for ct, (on, off) in enumerate(onsets_offsets):
        i = np.flatnonzero(timeline > on)[0]
        j = np.flatnonzero(timeline < off)[-1]
        out[ct] = [i, j]
    return out

def dictlist_to_arraydict(dictlist):
    """Convert from pupil format list of dicts to dict of arrays"""
    dict_fields = list(dictlist[0].keys())
    out = {}
    for df in dict_fields:
        out[df] = np.array([d[df] for d in dictlist])
    return out

def arraydict_to_dictlist(arraydict):
    """Convert from dict of arrays to pupil format list of dicts"""
    dict_fields = list(arraydict.keys())
    first_key = dict_fields[0]
    n = len(arraydict[first_key])
    out = []
    for j in range(n):
        frame_dict = {}
        for k in dict_fields:
            value = arraydict[k][j]
            if isinstance(value, np.ndarray):
                value = value.tolist()
            frame_dict[k] = value
        out.append(frame_dict)
    return out

def read_pl_gaze_csv(session_folder, output_id):
    sub_directory = str(output_id) * 3
    csv_file_name = os.path.join(session_folder,'exports',sub_directory,"gaze_positions.csv")
    print("CSV File Name: ", csv_file_name)
    return pd.read_csv(csv_file_name)


text:
"""Tease out some differences."""
import glob

from tqdm import tqdm
import numpy as np
import pyiem.dep as dep_utils
from pyiem.plot.use_agg import plt
import pandas as pd

MYHUCS = [
    "070600060701",
    "070801020905",
    "101702040703",
    "070600020602",
    "071000090201",
    "102300020202",
    "102300060407",
    "071000091101",
    "101702032003",
    "071000081503",
    "102300031404",
    "070801021001",
    "101702031903",
    "071000070402",
    "102400030501",
    "070600040401",
    "102300030402",
    "070802050304",
    "070801060604",
    "071000061401",
    "102802010402",
    "102400090404",
    "070802070602",
    "071000050703",
    "070802090401",
    "102400050501",
    "070802060403",
]

def main():
    """Go Main Go."""
    fig, ax = plt.subplots(1, 1)

    for huc12 in tqdm(MYHUCS):
        dfs = []

        for scenario in [0, 1002]:
            for fn in glob.glob(
                "/i/%s/env/%s/%s/*.env" % (scenario, huc12[:8], huc12[8:])
            ):
                fpath = int(fn.split("/")[-1].split(".")[0].split("_")[1])
                df = dep_utils.read_env(fn)
                df["fpath"] = fpath
                df["av_det"] = df["av_det"] * 4.463
                df["scenario"] = scenario
                dfs.append(df)
        df = pd.concat(dfs)
        df["year"] = df["date"].dt.year
        gdf = (
            df[["year", "scenario", "fpath", "av_det"]]
            .groupby(["year", "scenario", "fpath"])
            .sum()
        )
        gdf = gdf.reset_index()
        ptiles = np.arange(95, 100.01, 0.1)
        pc1 = np.percentile(gdf[gdf["scenario"] == 0]["av_det"].values, ptiles)
        pc2 = np.percentile(
            gdf[gdf["scenario"] == 1002]["av_det"].values, ptiles
        )
        ax.plot(pc1, pc2)
        ax.scatter(pc1[-1], pc2[-1], s=40, marker="s", color="b", zorder=3)
    ax.grid(True)
    ax.set_title(
        "95th Percentile and Higher Yearly Soil Delivery\n"
        "Max value shown by dot"
    )
    ax.set_xlabel("Production Yearly Soil Delivery T/a")
    ax.set_ylabel("New Flowpaths Yearly Soil Delivery T/a")
    ymax = max([ax.get_xlim()[1], ax.get_ylim()[1]])
    ax.plot([0, ymax], [0, ymax], color="k", lw=2, label="x=y")
    fig.savefig("/tmp/test.png")

if __name__ == "__main__":
    main()

ext: py

max_forks_repo_name: akrherz/idep

reason: import numpy


text:
import numpy as np
import pandas as pd
import math
import matplotlib.pyplot as plt
import pywt

def features_estimation(signal, channel_name, fs, frame, step, plot=True):
    """
    Compute time, frequency and time-frequency features from signal.
    :param signal: numpy array signal.
    :param channel_name: string variable with the EMG channel name in analysis.
    :param fs: int variable with the sampling frequency used to acquire the signal
    :param frame: sliding window size
    :param step: sliding window step size
    :param plot: boolean variable to plot estimated features.
    :return: total_feature_matrix -- python Data-frame with.
    :return: features_names -- python list with
    """

    features_names = ['VAR', 'RMS', 'IEMG', 'MAV', 'LOG', 'WL', 'ACC', 'DASDV', 'ZC', 'WAMP', 'MYOP', "FR", "MNP", "TP",
                      "MNF", "MDF", "PKF", "WENT"]

    time_matrix = time_features_estimation(signal, frame, step)
    frequency_matrix = frequency_features_estimation(signal, fs, frame, step)
    time_frequency_matrix = time_frequency_features_estimation(signal, frame, step)
    total_feature_matrix = pd.DataFrame(np.column_stack((time_matrix, frequency_matrix, time_frequency_matrix)).T,
                                        index=features_names)

    print('EMG features were from channel {} extracted successfully'.format(channel_name))

    if plot:
        plot_features(signal, channel_name, fs, total_feature_matrix, step)

    return total_feature_matrix, features_names

def time_features_estimation(signal, frame, step):
    """
    Compute time features from signal using sliding window method.
    :param signal: numpy array signal.
    :param frame: sliding window size.
    :param step: sliding window step size.
    :return: time_features_matrix: narray matrix with the time features stacked by columns.
    """

    variance = []
    rms = []
    iemg = []
    mav = []
    log_detector = []
    wl = []
    aac = []
    dasdv = []
    zc = []
    wamp = []
    myop = []

    th = np.mean(signal) + 3 * np.std(signal)

    for i in range(frame, signal.size, step):
        x = signal[i - frame:i]

        variance.append(np.var(x))
        rms.append(np.sqrt(np.mean(x ** 2)))
        iemg.append(np.sum(abs(x)))  # Integral
        mav.append(np.sum(np.absolute(x)) / frame)  # Mean Absolute Value
        log_detector.append(np.exp(np.sum(np.log10(np.absolute(x))) / frame))
        wl.append(np.sum(abs(np.diff(x))))  # Wavelength
        aac.append(np.sum(abs(np.diff(x))) / frame)  # Average Amplitude Change
        dasdv.append(
            math.sqrt((1 / (frame - 1)) * np.sum((np.diff(x)) ** 2)))  # Difference absolute standard deviation value
        zc.append(zcruce(x, th))  # Zero-Crossing
        wamp.append(wilson_amplitude(x, th))  # Willison amplitude
        myop.append(myopulse(x, th))  # Myopulse percentage rate

    time_features_matrix = np.column_stack((variance, rms, iemg, mav, log_detector, wl, aac, dasdv, zc, wamp, myop))
    return time_features_matrix

def frequency_features_estimation(signal, fs, frame, step):
    """
    Compute frequency features from signal using sliding window method.
    :param signal: numpy array signal.
    :param fs: sampling frequency of the signal.
    :param frame: sliding window size
    :param step: sliding window step size
    :return: frequency_features_matrix: narray matrix with the frequency features stacked by columns.
    """

    fr = []
    mnp = []
    tot = []
    mnf = []
    mdf = []
    pkf = []

    for i in range(frame, signal.size, step):
        x = signal[i - frame:i]
        frequency, power = spectrum(x, fs)

        fr.append(frequency_ratio(frequency, power))  # Frequency ratio
        mnp.append(np.sum(power) / len(power))  # Mean power
        tot.append(np.sum(power))  # Total power
        mnf.append(mean_freq(frequency, power))  # Mean frequency
        mdf.append(median_freq(frequency, power))  # Median frequency
        pkf.append(frequency[power.argmax()])  # Peak frequency

    frequency_features_matrix = np.column_stack((fr, mnp, tot, mnf, mdf, pkf))

    return frequency_features_matrix

def time_frequency_features_estimation(signal, frame, step):
    """
    Compute time-frequency features from signal using sliding window method.
    :param signal: numpy array signal.
    :param frame: sliding window size
    :param step: sliding window step size
    :return: h_wavelet: list
    """
    h_wavelet = []

    for i in range(frame, signal.size, step):
        x = signal[i - frame:i]

        E_a, E = wavelet_energy(x, 'db2', 4)
        E.insert(0, E_a)
        E = np.asarray(E) / 100

        h_wavelet.append(-np.sum(E * np.log2(E)))

    return h_wavelet

def wilson_amplitude(signal, th):
    x = abs(np.diff(signal))
    umbral = x >= th
    return np.sum(umbral)

def myopulse(signal, th):
    umbral = signal >= th
    return np.sum(umbral) / len(signal)

def spectrum(signal, fs):
    m = len(signal)
    n = next_power_of_2(m)
    y = np.fft.fft(signal, n)
    yh = y[0:int(n / 2 - 1)]
    fh = (fs / n) * np.arange(0, n / 2 - 1, 1)
    power = np.real(yh * np.conj(yh) / n)

    return fh, power

def frequency_ratio(frequency, power):
    power_low = power[(frequency >= 30) & (frequency <= 250)]
    power_high = power[(frequency > 250) & (frequency <= 500)]
    ULC = np.sum(power_low)
    UHC = np.sum(power_high)

    return ULC / UHC

def shannon(x):
    n = len(x)
    nb = 19
    hist, bin_edges = np.histogram(x, bins=nb)
    counts = hist / n
    nz = np.nonzero(counts)

    return np.sum(counts[nz] * np.log(counts[nz]) / np.log(2))

def zcruce(X, th):
    th = 0
    cruce = 0
    for cont in range(len(X) - 1):
        can = X[cont] * X[cont + 1]
        can2 = abs(X[cont] - X[cont + 1])
        if can < 0 and can2 > th:
            cruce = cruce + 1
    return cruce

def mean_freq(frequency, power):
    num = 0
    den = 0
    for i in range(int(len(power) / 2)):
        num += frequency[i] * power[i]
        den += power[i]

    return num / den

def median_freq(frequency, power):
    power_total = np.sum(power) / 2
    temp = 0
    tol = 0.01
    errel = 1
    i = 0

    while abs(errel) > tol:
        temp += power[i]
        errel = (power_total - temp) / power_total
        i += 1
        if errel < 0:
            errel = 0
            i -= 1

    return frequency[i]

def wavelet_energy(x, mother, nivel):
    coeffs = pywt.wavedecn(x, wavelet=mother, level=nivel)
    arr, _ = pywt.coeffs_to_array(coeffs)
    et = np.sum(arr ** 2)
    ca = coeffs[0]
    ea = 100 * np.sum(ca ** 2) / et
    ed = []

    for k in range(1, len(coeffs)):
        cd = list(coeffs[k].values())
        cd = np.asarray(cd)
        ed.append(100 * np.sum(cd ** 2) / et)

    return ea, ed

def next_power_of_2(x):
    return 1 if x == 0 else 2 ** (x - 1).bit_length()

def med_freq(f, P):
    plot = np.sum(P) / 2
    temp = 0
    tol = 0.01
    errel = 1
    i = 0

    while abs(errel) > tol:
        temp += P[i]
        errel = (plot - temp) / plot
        i += 1
        if errel < 0:
            errel = 0
            i -= 1

    return f[i]

def plot_features(signal, channel_name, fs, feature_matrix, step):
    """
    Argument:
    signal -- python numpy array representing recording of a signal.
    channel_name -- string variable with the EMG channel name in analysis (Title).
    fs -- int variable with the sampling frequency used to acquire the signal.
    feature_matrix -- python Dataframe.
    step -- int variable with the step size used in the sliding window method.
    """

    ts = np.arange(0, len(signal) / fs, 1 / fs)
    # for idx, f in enumerate(tfeatures.T):
    for key in feature_matrix.T:
        tf = step * (np.arange(0, len(feature_matrix.T[key]) / fs, 1 / fs))
        fig = plt.figure()

        ax = fig.add_subplot(111, label="1")
        ax2 = fig.add_subplot(111, label="2", frame_on=False)
        ax.plot(ts, signal, color="C0")
        ax.autoscale(tight=True)
        plt.title(channel_name + ": " + key)
        ax.set_xlabel("Time")
        ax.set_ylabel("mV")

        ax2.plot(tf, feature_matrix.T[key], color="red")
        ax2.yaxis.tick_right()
        ax2.autoscale(tight=True)
        ax2.set_xticks([])
        ax2.set_yticks([])
        # mng = plt.get_current_fig_manager()
        # mng.window.state('zoomed')
        plt.show()


text:
# pylint: disable=missing-docstring
import unittest

import tensorflow as tf
import numpy as np
import pytest

import tf_encrypted as tfe
from tf_encrypted.layers.pooling import MaxPooling2D

@pytest.mark.slow
class TestMaxPooling2D(unittest.TestCase):
  def setUp(self):
    tf.reset_default_graph()

  def tearDown(self):
    tf.reset_default_graph()

  def test_maxpool2d(self):
    with tfe.protocol.SecureNN() as prot:

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

      expected = np.array([[[[3, 4],
                             [3, 4]]]], dtype=np.float64)

      x = prot.define_private_variable(x_in)
      pool = MaxPooling2D([0, 1, 4, 4], pool_size=2, padding="VALID")
      result = pool.forward(x)

      with tfe.Session() as sess:
        sess.run(tf.global_variables_initializer())
        answer = sess.run(result.reveal())

    assert np.array_equal(answer, expected)

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

max_stars_repo_name: gavinuhma/tf-encrypted

max_issues_repo_name: dropoutlabs/tf-encrypted

max_forks_repo_name: dropoutlabs/tf-encrypted


text:
from simulation_game.buyer import MyopicBuyer, FaultyBuyer
from simulation_game.seller import DummySeller
from simulation_game.simulation import simulate

import numpy as np
import numpy.random as rn

price_scale = 3
mc_trials = 100
x_0 = 3
horizon = 4

def test_simulation_game_with_two_teams_same_strategies():
    rn.seed(123)
    buyer1 = MyopicBuyer("Sarah")
    seller1 = DummySeller("Sarah")
    buyer2 = MyopicBuyer("Steve")
    seller2 = DummySeller("Steve")
    teams = [(buyer1, seller1), (buyer2, seller2)]
    mean_revenue, mean_cs = simulate(teams, horizon, x_0, mc_trials, price_scale)
    assert mean_revenue[0] == mean_revenue[1] and mean_revenue[0] > 0.
    assert mean_cs[0] == mean_cs[1] and mean_cs[0] > 0.

def test_simulation_game_with_two_teams_different_strategies():
    rn.seed(1234)
    buyer1 = MyopicBuyer("Sarah")
    seller1 = DummySeller("Sarah")
    buyer2 = MyopicBuyer("Steve")
    seller2 = DummySeller("Steve", constant_price=0.1)
    teams = [(buyer1, seller1), (buyer2, seller2)]
    mean_revenue, mean_cs = simulate(teams, horizon, x_0, mc_trials, price_scale)
    # with a very low selling price for team 2, the seller is undervaluing the product and hence makes less revenue
    assert mean_revenue[0] > mean_revenue[1] and np.all(mean_revenue > 0.)
    # with a very low selling price for team 2, the buyer is making a good deal a lot of the time and has higher
    # consumer surplus, on average
    assert mean_cs[0] > mean_cs[1] and np.all(mean_cs > 0.)

def test_simulation_game_with_three_teams():
    rn.seed(1234)
    buyer1 = MyopicBuyer("Sarah")
    seller1 = DummySeller("Sarah")
    buyer2 = MyopicBuyer("Steve")
    seller2 = DummySeller("Steve", constant_price=0.1)
    buyer3 = MyopicBuyer("Emanuel")
    seller3 = DummySeller("Emanuel", constant_price=1.)
    teams = [(buyer1, seller1), (buyer2, seller2), (buyer3, seller3)]
    mean_revenue, mean_cs = simulate(teams, horizon, x_0, mc_trials, price_scale)
    # with a very low selling price for team 2, the seller is undervaluing the product and hence makes less revenue
    # teams 1 and 3 have the same selling strategy, so should get the same sample mean revenue
    assert mean_revenue[2] == mean_revenue[0] and mean_revenue[0] > mean_revenue[1]
    assert np.all(mean_revenue > 0.)
    # with a very low selling price for team 2, the buyer is making a good deal a lot of the time and has higher
    # consumer surplus, on average
    assert mean_cs[0] > mean_cs[1]
    assert np.all(mean_cs > 0.)

def test_faulty_bots_scenario():
    rn.seed(1234)
    buyer1 = MyopicBuyer("Sarah")
    seller1 = DummySeller("Sarah")
    buyer2 = FaultyBuyer()
    seller2 = DummySeller("Basil")
    teams = [(buyer1, seller1), (buyer2, seller2)]
    # This code should terminate.. that's all we're testing
    simulate(teams, horizon, x_0, 1, price_scale)

test_simulation_game_with_two_teams_same_strategies()
test_simulation_game_with_two_teams_different_strategies()
test_simulation_game_with_three_teams()
test_faulty_bots_scenario()

max_stars_repo_name: ecpaulson/OM-Game

max_issues_repo_name: ecpaulson/OM-Game

max_forks_repo_name: ecpaulson/OM-Game


text:
from PIL import Image
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import torch
import numpy as np
import cv2
import time
from samples.CLS2IDX import CLS2IDX
#@title Imagenet class indices to names

from baselines.ViT.ViT_LRP import vit_base_patch16_224 as vit_LRP #获得一个预训练模型
from baselines.ViT.ViT_explanation_generator import LRP #获得一个用于可解释性的生成器

normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    normalize,
])#图像处理

# create heatmap from mask on image
def show_cam_on_image(img, mask):
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    return cam

# initialize ViT pretrained
model = vit_LRP(pretrained=True).cuda() #预训练
model.eval() #测试模式
attribution_generator = LRP(model) #! 生成可解释模型

def generate_visualization(original_image, class_index=None):
    transformer_attribution = attribution_generator.generate_LRP(original_image.unsqueeze(0).cuda(), method="transformer_attribution", index=class_index).detach()
    transformer_attribution = transformer_attribution.reshape(1, 1, 14, 14)
    transformer_attribution = torch.nn.functional.interpolate(transformer_attribution, scale_factor=16, mode='bilinear')
    transformer_attribution = transformer_attribution.reshape(224, 224).cuda().data.cpu().numpy()
    transformer_attribution = (transformer_attribution - transformer_attribution.min()) / (transformer_attribution.max() - transformer_attribution.min())
    image_transformer_attribution = original_image.permute(1, 2, 0).data.cpu().numpy()
    image_transformer_attribution = (image_transformer_attribution - image_transformer_attribution.min()) / (image_transformer_attribution.max() - image_transformer_attribution.min())
    vis = show_cam_on_image(image_transformer_attribution, transformer_attribution)
    vis =  np.uint8(255 * vis)
    vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR)
    return vis

def print_top_classes(predictions, **kwargs):
    # Print Top-5 predictions
    prob = torch.softmax(predictions, dim=1)
    class_indices = predictions.data.topk(5, dim=1)[1][0].tolist()
    max_str_len = 0
    class_names = []
    for cls_idx in class_indices:
        class_names.append(CLS2IDX[cls_idx])
        if len(CLS2IDX[cls_idx]) > max_str_len:
            max_str_len = len(CLS2IDX[cls_idx])

    print('Top 5 classes:')
    for cls_idx in class_indices:
        output_string = '\t{} : {}'.format(cls_idx, CLS2IDX[cls_idx])
        output_string += ' ' * (max_str_len - len(CLS2IDX[cls_idx])) + '\t\t'
        output_string += 'value = {:.3f}\t prob = {:.1f}%'.format(predictions[0, cls_idx], 100 * prob[0, cls_idx])
        print(output_string)

def ViT_examples():
  image = Image.open('samples/el1.png')
  dog_cat_image = transform(image)

  fig, axs = plt.subplots(1, 3)
  axs[0].imshow(image);
  axs[0].axis('off');

  output = model(dog_cat_image.unsqueeze(0).cuda())
  print_top_classes(output)

  # cat - the predicted class
  cat = generate_visualization(dog_cat_image)

  # dog
  # generate visualization for class 243: 'bull mastiff'
  dog = generate_visualization(dog_cat_image, class_index=243)

  axs[1].imshow(cat)
  axs[1].axis('off')
  axs[2].imshow(dog)
  axs[2].axis('off')
  plt.savefig(f"./samples_output/{time.strftime('%Y-%m-%d %H:%M:%S')}.jpg")

ViT_examples()

max_stars_repo_name: TzRain/Transformer-Explainability

max_issues_repo_name: TzRain/Transformer-Explainability

max_forks_repo_name: TzRain/Transformer-Explainability


text:
import sys

import keras
import numpy
import pandas
from keras import layers
from keras.models import load_model

def assign_values_from_arguments():
    inp = ""
    mod = ""

    if len(sys.argv) != 5:
        print("Please re-run with correct arguments.")
        sys.exit()

    for i in range(len(sys.argv)):
        if sys.argv[i] == "-i":
            inp = sys.argv[i+1]
        elif sys.argv[i] == "-m":
            mod = sys.argv[i + 1]
    return inp, mod

if __name__ == '__main__':
    """
    The main function called when predict.py is run from the command line

    >Execute: python3.7 new_representation.py -i path_to:nn_representations.csv -m path_to:WindDenseNN.h5
    """
    inputFile, inputModel = assign_values_from_arguments()
    weights = load_model(inputModel).layers[0].get_weights()
    # If we uncomment the following line, we notice that relu is used as activation to this layer
    # print(load_model(inputModel).layers[0].output)
    model = keras.Sequential()
    model.add(layers.Dense(64, activation='relu', input_shape=(128,)))
    model.set_weights(weights)

    data = pandas.read_csv(inputFile, header=None)
    timestamps = data.loc[:, 0]
    for index, each in enumerate(timestamps):
        timestamps[index] = each.replace(" ", ",")
    data = data.drop(data.columns[0], axis=1)

    # model.summary()

    result = numpy.c_[timestamps.to_numpy(), model.predict(data, batch_size=32)]
    beginning_of_file = 'vectors'
    numpy.savetxt("new_representation.csv", result, delimiter="\t", fmt='%s', header=beginning_of_file, comments='')


text:
# coding=utf-8
# Copyright 2018 the HuggingFace Inc. team.
#
# 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 dataclasses
import gc
import os
import tempfile
import unittest

import numpy as np

from transformers import AutoTokenizer, IntervalStrategy, PretrainedConfig, TrainingArguments, is_torch_available
from transformers.file_utils import WEIGHTS_NAME
from transformers.testing_utils import (
    TestCasePlus,
    get_tests_dir,
    require_datasets,
    require_optuna,
    require_ray,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.utils.hp_naming import TrialShortNamer

if is_torch_available():
    import torch
    from torch.utils.data import IterableDataset

    from transformers import (
        AutoModelForMaskedLM,
        AutoModelForSequenceClassification,
        DataCollatorForLanguageModeling,
        EarlyStoppingCallback,
        GlueDataset,
        GlueDataTrainingArguments,
        GPT2Config,
        GPT2LMHeadModel,
        LineByLineTextDataset,
        PreTrainedModel,
        TextDataset,
        Trainer,
        TrainerState,
    )
    from transformers.modeling_utils import unwrap_model

PATH_SAMPLE_TEXT = f"{get_tests_dir()}/fixtures/sample_text.txt"

class RegressionDataset:
    def __init__(self, a=2, b=3, length=64, seed=42, label_names=None):
        np.random.seed(seed)
        self.label_names = ["labels"] if label_names is None else label_names
        self.length = length
        self.x = np.random.normal(size=(length,)).astype(np.float32)
        self.ys = [a * self.x + b + np.random.normal(scale=0.1, size=(length,)) for _ in self.label_names]
        self.ys = [y.astype(np.float32) for y in self.ys]

    def __len__(self):
        return self.length

    def __getitem__(self, i):
        result = {name: y[i] for name, y in zip(self.label_names, self.ys)}
        result["input_x"] = self.x[i]
        return result

@dataclasses.dataclass
class RegressionTrainingArguments(TrainingArguments):
    a: float = 0.0
    b: float = 0.0

class RepeatDataset:
    def __init__(self, x, length=64):
        self.x = x
        self.length = length

    def __len__(self):
        return self.length

    def __getitem__(self, i):
        return {"input_ids": self.x, "labels": self.x}

class DynamicShapesDataset:
    def __init__(self, length=64, seed=42, batch_size=8):
        self.length = length
        np.random.seed(seed)
        sizes = np.random.randint(1, 20, (length // batch_size,))
        # For easy batching, we make every batch_size consecutive samples the same size.
        self.xs = [np.random.normal(size=(s,)) for s in sizes.repeat(batch_size)]
        self.ys = [np.random.normal(size=(s,)) for s in sizes.repeat(batch_size)]

    def __len__(self):
        return self.length

    def __getitem__(self, i):
        return {"input_x": self.xs[i], "labels": self.ys[i]}

class AlmostAccuracy:
    def __init__(self, thresh=0.25):
        self.thresh = thresh

    def __call__(self, eval_pred):
        predictions, labels = eval_pred
        true = np.abs(predictions - labels) <= self.thresh
        return {"accuracy": true.astype(np.float32).mean().item()}

class RegressionModelConfig(PretrainedConfig):
    def __init__(self, a=0, b=0, double_output=False, **kwargs):
        super().__init__(**kwargs)
        self.a = a
        self.b = b
        self.double_output = double_output

if is_torch_available():

    class SampleIterableDataset(IterableDataset):
        """
        Criteria is not whether it is IterableDataset or not, criteria is whether __len__ is implemented
        """

        def __init__(self, file_path, tokenizer):
            self.ds = TextDataset(file_path=file_path, tokenizer=tokenizer, block_size=64)

        def __iter__(self):
            for i in range(len(self.ds)):
                yield self.ds[i]

    class RegressionModel(torch.nn.Module):
        def __init__(self, a=0, b=0, double_output=False):
            super().__init__()
            self.a = torch.nn.Parameter(torch.tensor(a).float())
            self.b = torch.nn.Parameter(torch.tensor(b).float())
            self.double_output = double_output
            self.config = None

        def forward(self, input_x, labels=None, **kwargs):
            y = input_x * self.a + self.b
            if labels is None:
                return (y, y) if self.double_output else (y,)
            loss = torch.nn.functional.mse_loss(y, labels)
            return (loss, y, y) if self.double_output else (loss, y)

    class RegressionDictModel(torch.nn.Module):
        def __init__(self, a=0, b=0):
            super().__init__()
            self.a = torch.nn.Parameter(torch.tensor(a).float())
            self.b = torch.nn.Parameter(torch.tensor(b).float())
            self.config = None

        def forward(self, input_x, labels=None, **kwargs):
            y = input_x * self.a + self.b
            result = {"output": y}
            if labels is not None:
                result["loss"] = torch.nn.functional.mse_loss(y, labels)
            return result

    class RegressionPreTrainedModel(PreTrainedModel):
        config_class = RegressionModelConfig
        base_model_prefix = "regression"

        def __init__(self, config):
            super().__init__(config)
            self.a = torch.nn.Parameter(torch.tensor(config.a).float())
            self.b = torch.nn.Parameter(torch.tensor(config.b).float())
            self.double_output = config.double_output

        def forward(self, input_x, labels=None, **kwargs):
            y = input_x * self.a + self.b
            if labels is None:
                return (y, y) if self.double_output else (y,)
            loss = torch.nn.functional.mse_loss(y, labels)
            return (loss, y, y) if self.double_output else (loss, y)

    class TstLayer(torch.nn.Module):
        def __init__(self, hidden_size):
            super().__init__()
            self.linear1 = torch.nn.Linear(hidden_size, hidden_size)
            self.ln1 = torch.nn.LayerNorm(hidden_size)
            self.linear2 = torch.nn.Linear(hidden_size, hidden_size)
            self.ln2 = torch.nn.LayerNorm(hidden_size)
            self.bias = torch.nn.Parameter(torch.zeros(hidden_size))

        def forward(self, x):
            h = self.ln1(torch.nn.functional.relu(self.linear1(x)))
            h = torch.nn.functional.relu(self.linear2(x))
            return self.ln2(x + h + self.bias)

    def get_regression_trainer(a=0, b=0, double_output=False, train_len=64, eval_len=64, pretrained=True, **kwargs):
        label_names = kwargs.get("label_names", None)
        train_dataset = RegressionDataset(length=train_len, label_names=label_names)
        eval_dataset = RegressionDataset(length=eval_len, label_names=label_names)
        if pretrained:
            config = RegressionModelConfig(a=a, b=b, double_output=double_output)
            model = RegressionPreTrainedModel(config)
        else:
            model = RegressionModel(a=a, b=b, double_output=double_output)
        compute_metrics = kwargs.pop("compute_metrics", None)
        data_collator = kwargs.pop("data_collator", None)
        optimizers = kwargs.pop("optimizers", (None, None))
        output_dir = kwargs.pop("output_dir", "./regression")
        model_init = kwargs.pop("model_init", None)

        args = RegressionTrainingArguments(output_dir, a=a, b=b, **kwargs)
        return Trainer(
            model,
            args,
            data_collator=data_collator,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            compute_metrics=compute_metrics,
            optimizers=optimizers,
            model_init=model_init,
        )

class TrainerIntegrationCommon:
    def check_saved_checkpoints(self, output_dir, freq, total, is_pretrained=True):
        file_list = [WEIGHTS_NAME, "training_args.bin", "optimizer.pt", "scheduler.pt", "trainer_state.json"]
        if is_pretrained:
            file_list.append("config.json")
        for step in range(freq, total, freq):
            checkpoint = os.path.join(output_dir, f"checkpoint-{step}")
            self.assertTrue(os.path.isdir(checkpoint))
            for filename in file_list:
                self.assertTrue(os.path.isfile(os.path.join(checkpoint, filename)))

    def check_best_model_has_been_loaded(
        self, output_dir, freq, total, trainer, metric, greater_is_better=False, is_pretrained=True
    ):
        checkpoint = os.path.join(output_dir, f"checkpoint-{(total // freq) * freq}")
        log_history = TrainerState.load_from_json(os.path.join(checkpoint, "trainer_state.json")).log_history

        values = [d[metric] for d in log_history]
        best_value = max(values) if greater_is_better else min(values)
        best_checkpoint = (values.index(best_value) + 1) * freq
        checkpoint = os.path.join(output_dir, f"checkpoint-{best_checkpoint}")
        if is_pretrained:
            best_model = RegressionPreTrainedModel.from_pretrained(checkpoint)
            best_model.to(trainer.args.device)
        else:
            best_model = RegressionModel()
            state_dict = torch.load(os.path.join(checkpoint, WEIGHTS_NAME))
            best_model.load_state_dict(state_dict)
            best_model.to(trainer.args.device)
        self.assertTrue(torch.allclose(best_model.a, trainer.model.a))
        self.assertTrue(torch.allclose(best_model.b, trainer.model.b))

        metrics = trainer.evaluate()
        self.assertEqual(metrics[metric], best_value)

    def check_trainer_state_are_the_same(self, trainer_state, trainer_state1):
        # We'll pop things so operate on copies.
        state = trainer_state.copy()
        state1 = trainer_state1.copy()
        # Log history main contain different logs for the time metrics (after resuming a training).
        log_history = state.pop("log_history", None)
        log_history1 = state1.pop("log_history", None)
        self.assertEqual(state, state1)
        for log, log1 in zip(log_history, log_history1):
            _ = log.pop("train_runtime", None)
            _ = log1.pop("train_runtime", None)
            _ = log.pop("train_samples_per_second", None)
            _ = log1.pop("train_samples_per_second", None)
            self.assertEqual(log, log1)

@require_torch
@require_sentencepiece
@require_tokenizers
class TrainerIntegrationTest(TestCasePlus, TrainerIntegrationCommon):
    def setUp(self):
        super().setUp()
        args = TrainingArguments(".")
        self.n_epochs = args.num_train_epochs
        self.batch_size = args.train_batch_size
        trainer = get_regression_trainer(learning_rate=0.1)
        trainer.train()
        self.default_trained_model = (trainer.model.a, trainer.model.b)

        trainer = get_regression_trainer(learning_rate=0.1, seed=314)
        trainer.train()
        self.alternate_trained_model = (trainer.model.a, trainer.model.b)

    def check_trained_model(self, model, alternate_seed=False):
        # Checks a training seeded with learning_rate = 0.1
        (a, b) = self.alternate_trained_model if alternate_seed else self.default_trained_model
        self.assertTrue(torch.allclose(model.a, a))
        self.assertTrue(torch.allclose(model.b, b))

    def test_trainer_works_with_dict(self):
        # Edge case because Apex with mode O2 will change our models to return dicts. This test checks it doesn't break
        # anything.
        train_dataset = RegressionDataset()
        eval_dataset = RegressionDataset()
        model = RegressionDictModel()
        args = TrainingArguments("./regression")
        trainer = Trainer(model, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
        trainer.train()
        _ = trainer.evaluate()
        _ = trainer.predict(eval_dataset)

    def test_evaluation_with_keys_to_drop(self):
        config = GPT2Config(vocab_size=100, n_positions=128, n_ctx=128, n_embd=32, n_layer=3, n_head=4)
        tiny_gpt2 = GPT2LMHeadModel(config)
        x = torch.randint(0, 100, (128,))
        eval_dataset = RepeatDataset(x)
        args = TrainingArguments("./test")
        trainer = Trainer(tiny_gpt2, args, eval_dataset=eval_dataset)
        # By default the past_key_values are removed
        result = trainer.predict(eval_dataset)
        self.assertTrue(isinstance(result.predictions, np.ndarray))
        # We can still get them by setting ignore_keys to []
        result = trainer.predict(eval_dataset, ignore_keys=[])
        self.assertTrue(isinstance(result.predictions, tuple))
        self.assertEqual(len(result.predictions), 2)

    def test_training_arguments_are_left_untouched(self):
        trainer = get_regression_trainer()
        trainer.train()
        args = TrainingArguments("./regression")
        dict1, dict2 = args.to_dict(), trainer.args.to_dict()
        for key in dict1.keys():
            # Logging dir can be slightly different as they default to something with the time.
            if key != "logging_dir":
                self.assertEqual(dict1[key], dict2[key])

    def test_reproducible_training(self):
        # Checks that training worked, model trained and seed made a reproducible training.
        trainer = get_regression_trainer(learning_rate=0.1)
        trainer.train()
        self.check_trained_model(trainer.model)

        # Checks that a different seed gets different (reproducible) results.
        trainer = get_regression_trainer(learning_rate=0.1, seed=314)
        trainer.train()
        self.check_trained_model(trainer.model, alternate_seed=True)

    def test_number_of_steps_in_training(self):
        # Regular training has n_epochs * len(train_dl) steps
        trainer = get_regression_trainer(learning_rate=0.1)
        train_output = trainer.train()
        self.assertEqual(train_output.global_step, self.n_epochs * 64 / self.batch_size)

        # Check passing num_train_epochs works (and a float version too):
        trainer = get_regression_trainer(learning_rate=0.1, num_train_epochs=1.5)
        train_output = trainer.train()
        self.assertEqual(train_output.global_step, int(1.5 * 64 / self.batch_size))

        # If we pass a max_steps, num_train_epochs is ignored
        trainer = get_regression_trainer(learning_rate=0.1, max_steps=10)
        train_output = trainer.train()
        self.assertEqual(train_output.global_step, 10)

    def test_train_and_eval_dataloaders(self):
        n_gpu = max(1, torch.cuda.device_count())
        trainer = get_regression_trainer(learning_rate=0.1, per_device_train_batch_size=16)
        self.assertEqual(trainer.get_train_dataloader().batch_size, 16 * n_gpu)
        trainer = get_regression_trainer(learning_rate=0.1, per_device_eval_batch_size=16)
        self.assertEqual(trainer.get_eval_dataloader().batch_size, 16 * n_gpu)

        # Check drop_last works
        trainer = get_regression_trainer(
            train_len=66, eval_len=74, learning_rate=0.1, per_device_train_batch_size=16, per_device_eval_batch_size=32
        )
        self.assertEqual(len(trainer.get_train_dataloader()), 66 // (16 * n_gpu) + 1)
        self.assertEqual(len(trainer.get_eval_dataloader()), 74 // (32 * n_gpu) + 1)

        trainer = get_regression_trainer(
            train_len=66,
            eval_len=74,
            learning_rate=0.1,
            per_device_train_batch_size=16,
            per_device_eval_batch_size=32,
            dataloader_drop_last=True,
        )
        self.assertEqual(len(trainer.get_train_dataloader()), 66 // (16 * n_gpu))
        self.assertEqual(len(trainer.get_eval_dataloader()), 74 // (32 * n_gpu))

        # Check passing a new dataset for evaluation works
        new_eval_dataset = RegressionDataset(length=128)
        self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // (32 * n_gpu))

    @require_torch_multi_gpu
    def test_data_is_not_parallelized_when_model_is_parallel(self):
        model = RegressionModel()
        # Make the Trainer believe it's a parallelized model
        model.is_parallelizable = True
        model.model_parallel = True
        args = TrainingArguments("./regression", per_device_train_batch_size=16, per_device_eval_batch_size=16)
        trainer = Trainer(model, args, train_dataset=RegressionDataset(), eval_dataset=RegressionDataset())
        # Check the Trainer was fooled
        self.assertTrue(trainer.is_model_parallel)
        self.assertEqual(trainer.args.n_gpu, 1)

        # The batch size of the training and evaluation dataloaders should be 16, not 16 * n_gpu
        self.assertEqual(trainer.get_train_dataloader().batch_size, 16)
        self.assertEqual(len(trainer.get_train_dataloader()), 64 // 16)
        self.assertEqual(trainer.get_eval_dataloader().batch_size, 16)
        self.assertEqual(len(trainer.get_eval_dataloader()), 64 // 16)

    def test_evaluate(self):
        trainer = get_regression_trainer(a=1.5, b=2.5, compute_metrics=AlmostAccuracy())
        results = trainer.evaluate()

        x, y = trainer.eval_dataset.x, trainer.eval_dataset.ys[0]
        pred = 1.5 * x + 2.5
        expected_loss = ((pred - y) ** 2).mean()
        self.assertAlmostEqual(results["eval_loss"], expected_loss)
        expected_acc = AlmostAccuracy()((pred, y))["accuracy"]
        self.assertAlmostEqual(results["eval_accuracy"], expected_acc)

        # With a number of elements not a round multiple of the batch size
        trainer = get_regression_trainer(a=1.5, b=2.5, eval_len=66, compute_metrics=AlmostAccuracy())
        results = trainer.evaluate()

        x, y = trainer.eval_dataset.x, trainer.eval_dataset.ys[0]
        pred = 1.5 * x + 2.5
        expected_loss = ((pred - y) ** 2).mean()
        self.assertAlmostEqual(results["eval_loss"], expected_loss)
        expected_acc = AlmostAccuracy()((pred, y))["accuracy"]
        self.assertAlmostEqual(results["eval_accuracy"], expected_acc)

    def test_predict(self):
        trainer = get_regression_trainer(a=1.5, b=2.5)
        preds = trainer.predict(trainer.eval_dataset).predictions
        x = trainer.eval_dataset.x
        self.assertTrue(np.allclose(preds, 1.5 * x + 2.5))

        # With a number of elements not a round multiple of the batch size
        trainer = get_regression_trainer(a=1.5, b=2.5, eval_len=66)
        preds = trainer.predict(trainer.eval_dataset).predictions
        x = trainer.eval_dataset.x
        self.assertTrue(np.allclose(preds, 1.5 * x + 2.5))

        # With more than one output of the model
        trainer = get_regression_trainer(a=1.5, b=2.5, double_output=True)
        preds = trainer.predict(trainer.eval_dataset).predictions
        x = trainer.eval_dataset.x
        self.assertTrue(len(preds), 2)
        self.assertTrue(np.allclose(preds[0], 1.5 * x + 2.5))
        self.assertTrue(np.allclose(preds[1], 1.5 * x + 2.5))

        # With more than one output/label of the model
        trainer = get_regression_trainer(a=1.5, b=2.5, double_output=True, label_names=["labels", "labels_2"])
        outputs = trainer.predict(trainer.eval_dataset)
        preds = outputs.predictions
        labels = outputs.label_ids
        x = trainer.eval_dataset.x
        self.assertTrue(len(preds), 2)
        self.assertTrue(np.allclose(preds[0], 1.5 * x + 2.5))
        self.assertTrue(np.allclose(preds[1], 1.5 * x + 2.5))
        self.assertTrue(np.array_equal(labels[0], trainer.eval_dataset.ys[0]))
        self.assertTrue(np.array_equal(labels[1], trainer.eval_dataset.ys[1]))

    def test_dynamic_shapes(self):
        eval_dataset = DynamicShapesDataset(batch_size=self.batch_size)
        model = RegressionModel(a=2, b=1)
        args = TrainingArguments("./regression")
        trainer = Trainer(model, args, eval_dataset=eval_dataset)

        # Check evaluation can run to completion
        _ = trainer.evaluate()

        # Check predictions
        preds = trainer.predict(eval_dataset)
        for expected, seen in zip(eval_dataset.ys, preds.label_ids):
            self.assertTrue(np.array_equal(expected, seen[: expected.shape[0]]))
            self.assertTrue(np.all(seen[expected.shape[0] :] == -100))

        for expected, seen in zip(eval_dataset.xs, preds.predictions):
            self.assertTrue(np.array_equal(2 * expected + 1, seen[: expected.shape[0]]))
            self.assertTrue(np.all(seen[expected.shape[0] :] == -100))

        # Same tests with eval accumulation
        args = TrainingArguments("./regression", eval_accumulation_steps=2)
        trainer = Trainer(model, args, eval_dataset=eval_dataset)

        # Check evaluation can run to completion
        _ = trainer.evaluate()

        # Check predictions
        preds = trainer.predict(eval_dataset)
        for expected, seen in zip(eval_dataset.ys, preds.label_ids):
            self.assertTrue(np.array_equal(expected, seen[: expected.shape[0]]))
            self.assertTrue(np.all(seen[expected.shape[0] :] == -100))

        for expected, seen in zip(eval_dataset.xs, preds.predictions):
            self.assertTrue(np.array_equal(2 * expected + 1, seen[: expected.shape[0]]))
            self.assertTrue(np.all(seen[expected.shape[0] :] == -100))

    @require_datasets
    def test_trainer_with_datasets(self):
        import datasets

        np.random.seed(42)
        x = np.random.normal(size=(64,)).astype(np.float32)
        y = 2.0 * x + 3.0 + np.random.normal(scale=0.1, size=(64,))
        train_dataset = datasets.Dataset.from_dict({"input_x": x, "label": y})

        # Base training. Should have the same results as test_reproducible_training
        model = RegressionModel()
        args = TrainingArguments("./regression", learning_rate=0.1)
        trainer = Trainer(model, args, train_dataset=train_dataset)
        trainer.train()
        self.check_trained_model(trainer.model)

        # Can return tensors.
        train_dataset.set_format(type="torch", dtype=torch.float32)
        model = RegressionModel()
        trainer = Trainer(model, args, train_dataset=train_dataset)
        trainer.train()
        self.check_trained_model(trainer.model)

        # Adding one column not used by the model should have no impact
        z = np.random.normal(size=(64,)).astype(np.float32)
        train_dataset = datasets.Dataset.from_dict({"input_x": x, "label": y, "extra": z})
        model = RegressionModel()
        trainer = Trainer(model, args, train_dataset=train_dataset)
        trainer.train()
        self.check_trained_model(trainer.model)

    def test_custom_optimizer(self):
        train_dataset = RegressionDataset()
        args = TrainingArguments("./regression")
        model = RegressionModel()
        optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
        lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: 1.0)
        trainer = Trainer(model, args, train_dataset=train_dataset, optimizers=(optimizer, lr_scheduler))
        trainer.train()

        (a, b) = self.default_trained_model
        self.assertFalse(torch.allclose(trainer.model.a, a))
        self.assertFalse(torch.allclose(trainer.model.b, b))
        self.assertEqual(trainer.optimizer.state_dict()["param_groups"][0]["lr"], 1.0)

    def test_model_init(self):
        train_dataset = RegressionDataset()
        args = TrainingArguments("./regression", learning_rate=0.1)
        trainer = Trainer(args=args, train_dataset=train_dataset, model_init=lambda: RegressionModel())
        trainer.train()
        self.check_trained_model(trainer.model)

        # Re-training should restart from scratch, thus lead the same results.
        trainer.train()
        self.check_trained_model(trainer.model)

        # Re-training should restart from scratch, thus lead the same results and new seed should be used.
        trainer.args.seed = 314
        trainer.train()
        self.check_trained_model(trainer.model, alternate_seed=True)

    def test_save_checkpoints(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(output_dir=tmpdir, save_steps=5)
            trainer.train()
            self.check_saved_checkpoints(tmpdir, 5, int(self.n_epochs * 64 / self.batch_size))

        # With a regular model that is not a PreTrainedModel
        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(output_dir=tmpdir, save_steps=5, pretrained=False)
            trainer.train()
            self.check_saved_checkpoints(tmpdir, 5, int(self.n_epochs * 64 / self.batch_size), False)

    def test_gradient_accumulation(self):
        # Training with half the batch size but accumulation steps as 2 should give the same results.
        trainer = get_regression_trainer(
            gradient_accumulation_steps=2, per_device_train_batch_size=4, learning_rate=0.1
        )
        trainer.train()
        self.check_trained_model(trainer.model)

    @require_torch_multi_gpu
    def test_run_seq2seq_double_train_wrap_once(self):
        # test that we don't wrap the model more than once
        # since wrapping primarily happens on multi-gpu setup we want multiple gpus to test for
        # example DataParallel(DataParallel(model))

        trainer = get_regression_trainer()
        trainer.train()
        model_wrapped_before = trainer.model_wrapped
        trainer.train()
        model_wrapped_after = trainer.model_wrapped
        self.assertIs(model_wrapped_before, model_wrapped_after, "should be not wrapped twice")

    def test_can_resume_training(self):
        if torch.cuda.device_count() > 2:
            # This test will fail for more than 2 GPUs since the batch size will get bigger and with the number of
            # save_steps, the checkpoint will resume training at epoch 2 or more (so the data seen by the model
            # won't be the same since the training dataloader is shuffled).
            return

        with tempfile.TemporaryDirectory() as tmpdir:
            kwargs = dict(output_dir=tmpdir, train_len=128, save_steps=5, learning_rate=0.1)
            trainer = get_regression_trainer(**kwargs)
            trainer.train()
            (a, b) = trainer.model.a.item(), trainer.model.b.item()
            state = dataclasses.asdict(trainer.state)

            checkpoint = os.path.join(tmpdir, "checkpoint-5")

            # Reinitialize trainer
            trainer = get_regression_trainer(**kwargs)

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

            # Now check with a later checkpoint that it also works when we span over one epoch
            checkpoint = os.path.join(tmpdir, "checkpoint-15")

            # Reinitialize trainer and load model
            trainer = get_regression_trainer(**kwargs)

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

        # With a regular model that is not a PreTrainedModel
        with tempfile.TemporaryDirectory() as tmpdir:
            kwargs = dict(output_dir=tmpdir, train_len=128, save_steps=5, learning_rate=0.1, pretrained=False)

            trainer = get_regression_trainer(**kwargs)
            trainer.train()
            (a, b) = trainer.model.a.item(), trainer.model.b.item()
            state = dataclasses.asdict(trainer.state)

            checkpoint = os.path.join(tmpdir, "checkpoint-5")

            # Reinitialize trainer and load model
            trainer = get_regression_trainer(**kwargs)

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

            # Now check with a later checkpoint that it also works when we span over one epoch
            checkpoint = os.path.join(tmpdir, "checkpoint-15")

            # Reinitialize trainer and load model
            trainer = get_regression_trainer(**kwargs)

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

        # Now check failures

        # 1. fail to find a bogus checkpoint
        trainer = get_regression_trainer()
        with self.assertRaises(Exception) as context:
            trainer.train(resume_from_checkpoint=f"{checkpoint}-bogus")
        self.assertTrue("Can't find a valid checkpoint at" in str(context.exception))

        # 2. fail to find any checkpoint - due a fresh output_dir
        output_dir2 = self.get_auto_remove_tmp_dir()
        trainer = get_regression_trainer(output_dir=output_dir2)
        with self.assertRaises(Exception) as context:
            trainer.train(resume_from_checkpoint=True)
        self.assertTrue("No valid checkpoint found in output directory" in str(context.exception))

    def test_resume_training_with_gradient_accumulation(self):
        if torch.cuda.device_count() > 2:
            # This test will fail for more than 2 GPUs since the batch size will get bigger and with the number of
            # save_steps, the checkpoint will resume training at epoch 2 or more (so the data seen by the model
            # won't be the same since the training dataloader is shuffled).
            return
        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(
                output_dir=tmpdir,
                train_len=128,
                gradient_accumulation_steps=2,
                per_device_train_batch_size=4,
                save_steps=5,
                learning_rate=0.1,
            )
            trainer.train()
            (a, b) = trainer.model.a.item(), trainer.model.b.item()
            state = dataclasses.asdict(trainer.state)

            checkpoint = os.path.join(tmpdir, "checkpoint-5")

            # Reinitialize trainer
            trainer = get_regression_trainer(
                output_dir=tmpdir,
                train_len=128,
                gradient_accumulation_steps=2,
                per_device_train_batch_size=4,
                save_steps=5,
                learning_rate=0.1,
            )

            trainer.train(resume_from_checkpoint=checkpoint)
            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
            state1 = dataclasses.asdict(trainer.state)
            self.assertEqual(a, a1)
            self.assertEqual(b, b1)
            self.check_trainer_state_are_the_same(state, state1)

    def test_load_best_model_at_end(self):
        total = int(self.n_epochs * 64 / self.batch_size)
        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(
                a=1.5,
                b=2.5,
                output_dir=tmpdir,
                learning_rate=0.1,
                eval_steps=5,
                evaluation_strategy="steps",
                load_best_model_at_end=True,
            )
            self.assertFalse(trainer.args.greater_is_better)
            trainer.train()
            self.check_saved_checkpoints(tmpdir, 5, total)
            self.check_best_model_has_been_loaded(tmpdir, 5, total, trainer, "eval_loss")

        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(
                a=1.5,
                b=2.5,
                output_dir=tmpdir,
                learning_rate=0.1,
                eval_steps=5,
                evaluation_strategy="steps",
                load_best_model_at_end=True,
                metric_for_best_model="accuracy",
                compute_metrics=AlmostAccuracy(),
            )
            self.assertTrue(trainer.args.greater_is_better)
            trainer.train()
            self.check_saved_checkpoints(tmpdir, 5, total)
            self.check_best_model_has_been_loaded(tmpdir, 5, total, trainer, "eval_accuracy", greater_is_better=True)

        # Save is done every eval regardless of the strategy
        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(
                a=1.5,
                b=2.5,
                output_dir=tmpdir,
                learning_rate=0.1,
                evaluation_strategy="epoch",
                load_best_model_at_end=True,
                metric_for_best_model="accuracy",
                compute_metrics=AlmostAccuracy(),
            )
            self.assertTrue(trainer.args.greater_is_better)
            trainer.train()
            self.check_saved_checkpoints(tmpdir, 64 // self.batch_size, total)
            self.check_best_model_has_been_loaded(
                tmpdir, 64 // self.batch_size, total, trainer, "eval_accuracy", greater_is_better=True
            )

        # Test this works with a non PreTrainedModel
        with tempfile.TemporaryDirectory() as tmpdir:
            trainer = get_regression_trainer(
                output_dir=tmpdir,
                learning_rate=0.1,
                eval_steps=5,
                evaluation_strategy="steps",
                load_best_model_at_end=True,
                pretrained=False,
            )
            self.assertFalse(trainer.args.greater_is_better)
            trainer.train()
            self.check_saved_checkpoints(tmpdir, 5, total, is_pretrained=False)
            self.check_best_model_has_been_loaded(tmpdir, 5, total, trainer, "eval_loss", is_pretrained=False)

    @slow
    def test_trainer_eval_mrpc(self):
        MODEL_ID = "bert-base-cased-finetuned-mrpc"
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
        data_args = GlueDataTrainingArguments(
            task_name="mrpc", data_dir=f"{get_tests_dir()}/fixtures/tests_samples/MRPC", overwrite_cache=True
        )
        eval_dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")

        training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
        trainer = Trainer(model=model, args=training_args, eval_dataset=eval_dataset)
        result = trainer.evaluate()
        self.assertLess(result["eval_loss"], 0.2)

    @slow
    def test_trainer_eval_lm(self):
        MODEL_ID = "distilroberta-base"
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        dataset = LineByLineTextDataset(
            tokenizer=tokenizer,
            file_path=PATH_SAMPLE_TEXT,
            block_size=tokenizer.max_len_single_sentence,
        )
        self.assertEqual(len(dataset), 31)

    def test_trainer_iterable_dataset(self):
        # Simulate Language Modeling with an IterableDataset, with no __len__ method
        # Pick-up a tiny model, so it works on CPU
        # See Issue #5990: https://github.com/huggingface/transformers/issues/5990
        MODEL_ID = "sshleifer/tiny-distilbert-base-cased"
        model = AutoModelForMaskedLM.from_pretrained(MODEL_ID)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        train_dataset = SampleIterableDataset(file_path=PATH_SAMPLE_TEXT, tokenizer=tokenizer)
        training_args = TrainingArguments(output_dir="./examples", no_cuda=True, max_steps=2)
        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)

        training_args = TrainingArguments(output_dir="./examples", no_cuda=True, max_steps=2)
        trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, data_collator=data_collator)
        trainer.train()

        loader = trainer.get_train_dataloader()
        self.assertIsInstance(loader, torch.utils.data.DataLoader)
        self.assertIsInstance(loader.sampler, torch.utils.data.dataloader._InfiniteConstantSampler)

        # Exception if giving iterable dataset and no max_steps
        with self.assertRaises(ValueError):
            training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
            _ = Trainer(model=model, args=training_args, train_dataset=train_dataset, data_collator=data_collator)

        # Exception if eval_dataset is iterable in __init__
        with self.assertRaises(ValueError):
            training_args = TrainingArguments(output_dir="./examples", no_cuda=True, max_steps=2)
            _ = Trainer(
                model=model,
                args=training_args,
                train_dataset=train_dataset,
                eval_dataset=train_dataset,
                data_collator=data_collator,
            )

        # Exception if predicting with iterable dataset
        with self.assertRaises(ValueError):
            training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
            trainer = Trainer(model=model, args=training_args, data_collator=data_collator)
            trainer.predict(train_dataset)

        # Exception if evaluating with iterable dataset
        with self.assertRaises(ValueError):
            training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
            trainer = Trainer(model=model, args=training_args, data_collator=data_collator)
            trainer.evaluate(train_dataset)

    def test_num_train_epochs_in_training(self):
        # len(train_dl) < gradient_accumulation_steps shouldn't give ``ZeroDivisionError`` when ``max_steps`` is given.
        # It should give 1 update step for each epoch.
        trainer = get_regression_trainer(
            max_steps=3, train_len=64, per_device_train_batch_size=16, gradient_accumulation_steps=5
        )
        train_output = trainer.train()
        self.assertEqual(train_output.global_step, 3)

        # Even ``max_steps`` is not specified, we still expect 1 update step for each epoch if
        # len(train_dl) < gradient_accumulation_steps.
        trainer = get_regression_trainer(train_len=64, per_device_train_batch_size=16, gradient_accumulation_steps=5)
        train_output = trainer.train()
        self.assertEqual(train_output.global_step, int(self.n_epochs))

    def test_early_stopping_callback(self):
        # early stopping stops training before num_training_epochs
        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = get_regression_trainer(
                output_dir=tmp_dir,
                num_train_epochs=20,
                gradient_accumulation_steps=1,
                per_device_train_batch_size=16,
                load_best_model_at_end=True,
                evaluation_strategy=IntervalStrategy.EPOCH,
                compute_metrics=AlmostAccuracy(),
                metric_for_best_model="accuracy",
            )
            trainer.add_callback(EarlyStoppingCallback(1, 0.0001))
            train_output = trainer.train()
            self.assertLess(train_output.global_step, 20 * 64 / 16)

        # Invalid inputs to trainer with early stopping callback result in assertion error
        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = get_regression_trainer(
                output_dir=tmp_dir,
                num_train_epochs=20,
                gradient_accumulation_steps=1,
                per_device_train_batch_size=16,
                evaluation_strategy=IntervalStrategy.EPOCH,
                compute_metrics=AlmostAccuracy(),
                metric_for_best_model="accuracy",
            )
            trainer.add_callback(EarlyStoppingCallback(1))
            self.assertEqual(trainer.state.global_step, 0)
            try:
                trainer.train()
            except AssertionError:
                self.assertEqual(trainer.state.global_step, 0)

    def test_flos_extraction(self):
        trainer = get_regression_trainer(learning_rate=0.1)

        def assert_flos_extraction(trainer, wrapped_model_to_check):
            self.assertEqual(trainer.model, unwrap_model(wrapped_model_to_check))
            self.assertGreaterEqual(getattr(unwrap_model(wrapped_model_to_check).config, "total_flos", 0), 0)

        # with plain model
        assert_flos_extraction(trainer, trainer.model)

        # with enforced DataParallel
        assert_flos_extraction(trainer, torch.nn.DataParallel(trainer.model))

        trainer.train()
        self.assertTrue(isinstance(trainer.state.total_flos, float))

    def check_mem_metrics(self, trainer, check_func):
        metrics = trainer.train().metrics
        check_func("init_mem_cpu_alloc_delta", metrics)
        check_func("train_mem_cpu_alloc_delta", metrics)
        if torch.cuda.device_count() > 0:
            check_func("init_mem_gpu_alloc_delta", metrics)
            check_func("train_mem_gpu_alloc_delta", metrics)

        metrics = trainer.evaluate()
        check_func("eval_mem_cpu_alloc_delta", metrics)
        if torch.cuda.device_count() > 0:
            check_func("eval_mem_gpu_alloc_delta", metrics)

        metrics = trainer.predict(RegressionDataset()).metrics
        check_func("test_mem_cpu_alloc_delta", metrics)
        if torch.cuda.device_count() > 0:
            check_func("test_mem_gpu_alloc_delta", metrics)

    def test_mem_metrics(self):

        # with mem metrics enabled
        trainer = get_regression_trainer()
        self.check_mem_metrics(trainer, self.assertIn)

        # with mem metrics disabled
        trainer = get_regression_trainer(skip_memory_metrics=True)
        self.check_mem_metrics(trainer, self.assertNotIn)

    @require_torch_gpu
    def test_fp16_full_eval(self):

        # this is a sensitive test so let's keep debugging printouts in place for quick diagnosis.
        # it's using pretty large safety margins, but small enough to detect broken functionality.
        debug = 0

        bs = 8
        # make the params somewhat big so that there will be enough RAM consumed to be able to
        # measure things. We should get about 64KB for a+b in fp32
        a = torch.ones(1000, bs) + 0.001
        b = torch.ones(1000, bs) - 0.001

        # 1. with mem metrics enabled
        trainer = get_regression_trainer(a=a, b=b, eval_len=16)
        metrics = trainer.evaluate()
        del trainer
        gc.collect()

        fp32_init = metrics["init_mem_gpu_alloc_delta"]
        fp32_eval = metrics["eval_mem_gpu_alloc_delta"]

        if debug:
            print(f"fp32_init {fp32_init}")
            print(f"fp32_eval {fp32_eval}")

        # here we expect the model to be preloaded in trainer.__init__ and consume around 64K gpu ram.
        # perfect world: fp32_init == 64<<10
        self.assertGreater(fp32_init, 59_000)
        # after eval should be no extra memory allocated - with a small margin (other than the peak
        # memory consumption for the forward calculation that gets recovered)
        # perfect world: fp32_eval == close to zero
        self.assertLess(fp32_eval, 5_000)

        # 2. with mem metrics disabled
        trainer = get_regression_trainer(a=a, b=b, eval_len=16, fp16_full_eval=True)
        metrics = trainer.evaluate()
        fp16_init = metrics["init_mem_gpu_alloc_delta"]
        fp16_eval = metrics["eval_mem_gpu_alloc_delta"]

        if debug:
            print(f"fp16_init {fp16_init}")
            print(f"fp16_eval {fp16_eval}")

        # here we expect the model to not be preloaded in trainer.__init__, so with a small margin it should be close to 0
        # perfect world: fp16_init == close to zero
        self.assertLess(fp16_init, 5_000)
        # here we put the model on device in eval and only `half()` of it, i.e. about 32K,(again we ignore the peak margin which gets returned back)
        # perfect world: fp32_init == 32<<10
        self.assertGreater(fp16_eval, 27_000)

        # 3. relative comparison fp32 vs full fp16
        # should be about half of fp16_init
        # perfect world: fp32_init/2 == fp16_eval
        self.assertAlmostEqual(fp16_eval, fp32_init / 2, delta=5_000)

    def test_no_wd_param_group(self):
        model = torch.nn.Sequential(TstLayer(128), torch.nn.ModuleList([TstLayer(128), TstLayer(128)]))
        trainer = Trainer(model=model)
        trainer.create_optimizer_and_scheduler(10)
        # fmt: off
        wd_names = ['0.linear1.weight', '0.linear2.weight', '1.0.linear1.weight', '1.0.linear2.weight', '1.1.linear1.weight', '1.1.linear2.weight']
        # fmt: on
        wd_params = [p for n, p in model.named_parameters() if n in wd_names]
        no_wd_params = [p for n, p in model.named_parameters() if n not in wd_names]
        self.assertListEqual(trainer.optimizer.param_groups[0]["params"], wd_params)
        self.assertListEqual(trainer.optimizer.param_groups[1]["params"], no_wd_params)

@require_torch
@require_optuna
class TrainerHyperParameterOptunaIntegrationTest(unittest.TestCase):
    def setUp(self):
        args = TrainingArguments(".")
        self.n_epochs = args.num_train_epochs
        self.batch_size = args.train_batch_size

    def test_hyperparameter_search(self):
        class MyTrialShortNamer(TrialShortNamer):
            DEFAULTS = {"a": 0, "b": 0}

        def hp_space(trial):
            return {}

        def model_init(trial):
            if trial is not None:
                a = trial.suggest_int("a", -4, 4)
                b = trial.suggest_int("b", -4, 4)
            else:
                a = 0
                b = 0
            config = RegressionModelConfig(a=a, b=b, double_output=False)

            return RegressionPreTrainedModel(config)

        def hp_name(trial):
            return MyTrialShortNamer.shortname(trial.params)

        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = get_regression_trainer(
                output_dir=tmp_dir,
                learning_rate=0.1,
                logging_steps=1,
                evaluation_strategy=IntervalStrategy.EPOCH,
                num_train_epochs=4,
                disable_tqdm=True,
                load_best_model_at_end=True,
                logging_dir="runs",
                run_name="test",
                model_init=model_init,
            )
            trainer.hyperparameter_search(direction="minimize", hp_space=hp_space, hp_name=hp_name, n_trials=4)

@require_torch
@require_ray
class TrainerHyperParameterRayIntegrationTest(unittest.TestCase):
    def setUp(self):
        args = TrainingArguments(".")
        self.n_epochs = args.num_train_epochs
        self.batch_size = args.train_batch_size

    def test_hyperparameter_search(self):
        class MyTrialShortNamer(TrialShortNamer):
            DEFAULTS = {"a": 0, "b": 0}

        def hp_space(trial):
            from ray import tune

            return {
                "a": tune.randint(-4, 4),
                "b": tune.randint(-4, 4),
            }

        def model_init(config):
            model_config = RegressionModelConfig(a=config["a"], b=config["b"], double_output=False)

            return RegressionPreTrainedModel(model_config)

        def hp_name(params):
            return MyTrialShortNamer.shortname(params)

        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = get_regression_trainer(
                output_dir=tmp_dir,
                learning_rate=0.1,
                logging_steps=1,
                evaluation_strategy=IntervalStrategy.EPOCH,
                num_train_epochs=4,
                disable_tqdm=True,
                load_best_model_at_end=True,
                logging_dir="runs",
                run_name="test",
                model_init=model_init,
            )
            trainer.hyperparameter_search(
                direction="minimize", hp_space=hp_space, hp_name=hp_name, backend="ray", n_trials=4
            )

max_stars_repo_name: JadeMaveric/transformers

max_issues_repo_name: JadeMaveric/transformers

max_forks_repo_name: JadeMaveric/transformers


text:
from PDFContentConverter import PDFContentConverter
import networkx as nx
import numpy as np
from GraphConverter.util import RectangleUtil, constants
from GraphConverter.document.DocumentMetaCharacteristics import DocumentMetaCharacteristics
from GraphConverter.merging.PDFTextBoxMerging import PDFTextBoxMerging

class GraphConverter(object):

    def __init__(self, pdf, merge_boxes=False, regress_parameters=False,
                 use_font=True, use_width=True, use_rect=True, use_horizontal_overlap=False,
                 use_vertical_overlap=False,
                 page_ratio_x=2, page_ratio_y=2, x_eps=2, y_eps=2, font_eps_h=1, font_eps_v=1,
                 width_pct_eps=.4, width_page_eps=.5):
        self.pdf = pdf
        self.graphs = None
        self.meta = None
        self.set_attributes(merge_boxes, regress_parameters,
                            use_font, use_width, use_rect, use_horizontal_overlap, use_vertical_overlap,
                            page_ratio_x, page_ratio_y, x_eps, y_eps, font_eps_h, font_eps_v,
                            width_pct_eps, width_page_eps)
        conv = PDFContentConverter(self.pdf).convert()
        self.loc_df = conv["content"]
        self.loc_df["is_loop"] = 0
        self.media_boxes = conv["media_boxes"]
        self.n = conv["page_count"]
        self.loc_df = PDFTextBoxMerging(data=self.loc_df,
                                        media_boxes=self.media_boxes).transform()
        if self.loc_df is not None:
            self.meta = DocumentMetaCharacteristics(self.loc_df, self.media_boxes,
                                                    optimize=regress_parameters)
            self.meta.generate_attributes()

    def set_attributes(self, merge_boxes, regress_parameters,
                       use_font, use_width, use_rect, use_horizontal_overlap, use_vertical_overlap,
                       page_ratio_x, page_ratio_y, x_eps, y_eps, font_eps_h, font_eps_v,
                       width_pct_eps, width_page_eps):
        """

        :param merge_boxes:
        :param regress_parameters:
        :param use_font:
        :param use_width:
        :param use_rect:
        :param use_horizontal_overlap:
        :param use_vertical_overlap:
        :param page_ratio_x:
        :param page_ratio_y:
        :param x_eps:
        :param y_eps:
        :param font_eps_h:
        :param font_eps_v:
        :param width_pct_eps:
        :param width_page_eps:
        :return:
        """
        constants.USE_FONT = use_font
        constants.USE_WIDTH = use_width
        constants.USE_RECT = use_rect
        constants.USE_HORIZONTAL_OVERLAP = use_horizontal_overlap
        constants.USE_VERTICAL_OVERLAP = use_vertical_overlap
        constants.MERGE_RECTANGLE_TEXT_BOXES = merge_boxes
        constants.USE_RECTANGLE_BOX_COORDINATES = merge_boxes
        constants.OPTIMIZE = regress_parameters
        constants.DEFAULT_PAGE_RATIO_X = page_ratio_x
        constants.DEFAULT_PAGE_RATIO_Y = page_ratio_y
        constants.DEFAULT_X_EPS = x_eps
        constants.DEFAULT_Y_EPS = y_eps
        constants.DEFAULT_X_PHYS = x_eps
        constants.DEFAULT_FONT_EPS_HORIZONTAL = font_eps_h
        constants.DEFAULT_FONT_EPS_VERTICAL = font_eps_v
        constants.DEFAULT_WIDTH_PCT = width_pct_eps
        constants.DEFAULT_PAGE_WIDTH_EPS = width_page_eps

    def convert(self):
        """

        :return:
        """
        if self.loc_df is None:
            return [], None

        self.loc_df.set_index("id", inplace=True)
        self.loc_df["cluster_label"] = 0

        graph_list = []
        for page in range(self.meta.num_pages):
            G = nx.MultiDiGraph()
            loc_df_page = self.loc_df[self.loc_df["page"] == page]
            G.add_nodes_from(list(loc_df_page.index))
            nodes = loc_df_page.to_dict("index")
            nx.set_node_attributes(G, nodes)
            G = self.add_edges(G, page)
            G = nx.disjoint_union(nx.MultiDiGraph(), G)
            G = self.label_loops(G)
            graph_list.append(G)
        self.graphs = graph_list
        return {
            "graphs": graph_list,
            "document_meta": self.meta
        }

    def add_edges(self, G, page=0):
        """

        :param G:
        :param page:
        :return:
        """
        nodes = dict(G.nodes(data=True))
        node_ids = list(nodes.keys())
        for i in range(len(node_ids) - 1):
            id1 = node_ids[i]
            node1 = nodes[id1]

            for j in range(i + 1, len(node_ids)):
                id2 = node_ids[j]
                node2 = nodes[id2]

                delta_x_min = min(abs(node2["x_0"] - node1["x_0"]),
                                  abs(node2["x_1"] - node1["x_1"]),
                                  abs(node2["pos_x"] - node1["pos_x"]))
                delta_y_min = min(abs(node2["y_0"] - node1["y_0"]),
                                  abs(node2["y_1"] - node1["y_1"]),
                                  abs(node2["pos_y"] - node1["pos_y"]))
                length_x_min = min(abs(node2["x_1"] - node1["x_0"]),
                                   abs(node2["x_0"] - node1["x_1"]))
                length_y_min = min(abs(node2["y_1"] - node1["y_0"]),
                                   abs(node2["y_0"] - node1["y_1"]))
                delta_f = abs(node2["font_size"] - node1["font_size"])

                # horizontal edges
                if self.fulfills_horizontal_edge_conditions(node1, node2, page,
                                                            delta_f, length_x_min, delta_y_min):
                    if node1["pos_x"] < node2["pos_x"]:
                        G = self.get_horizontal_edge(G, id1, id2, length_x_min)
                    else:
                        G = self.get_horizontal_edge(G, id2, id1, length_x_min)

                # vertical edges
                if self.fulfills_vertical_edge_conditions(node1, node2, page,
                                                          delta_x_min, length_y_min, delta_f):
                    if node2["y_1"] > node1["y_1"]:
                        G = self.get_vertical_edge(G, id1, id2, length_y_min)
                    else:
                        G = self.get_vertical_edge(G, id2, id1, length_y_min)
        return G

    def get_horizontal_edge(self, G, id_start, id_end, length):
        """

        :param G:
        :param id_start:
        :param id_end:
        :param length:
        :return:
        """
        edges_from_node = [e for e in G.out_edges(id_start, keys=True, data=True)
                           if e[3]["direction"] == "h"]
        if len(edges_from_node) == 0:
            G.add_edge(id_start, id_end, length=length,
                       lengthx_phys=length, lengthy_phys=0, direction="h",
                       weight=length, key="h")
        elif len(edges_from_node) == 1:
            if edges_from_node[0][3]["lengthx_phys"] >= length:
                G.remove_edges_from(edges_from_node)
                G.add_edge(id_start, id_end, length=length,
                           lengthx_phys=length, lengthy_phys=0, direction="h",
                           weight=length, key="h")
        return G

    def get_vertical_edge(self, G, id_start, id_end, length):
        """

        :param G:
        :param id_start:
        :param id_end:
        :param length:
        :return:
        """
        edges_from_node_phys = [e for e in G.out_edges(id_start, keys=True, data=True)
                                if e[3]["direction"] == "v"]
        if len(edges_from_node_phys) == 0:
            G.add_edge(id_start, id_end, length=length,
                       lengthx_phys=0, lengthy_phys=length,
                       direction="v", weight=length, key="v")
        elif len(edges_from_node_phys) == 1:
            if edges_from_node_phys[0][3]["lengthy_phys"] >= length:
                G.remove_edges_from(edges_from_node_phys)
                G.add_edge(id_start, id_end, length=length,
                           lengthx_phys=0, lengthy_phys=length,
                           direction="v", weight=length, key="v")
        return G

    def label_loops(self, graph):
        """

        :param graph:
        :return:
        """
        loop_labels = dict()
        all_nodes = dict(graph.nodes(data=True))
        subgraph_H = nx.DiGraph()
        subgraph_V = nx.DiGraph()
        subgraph_H.add_nodes_from(graph)
        subgraph_V.add_nodes_from(graph)
        subgraph_H.add_edges_from([(e[0], e[1]) for e in graph.edges(keys=True, data=True) if e[2] == "h"])
        subgraph_V.add_edges_from([(e[0], e[1]) for e in graph.edges(keys=True, data=True) if e[2] == "v"])
        distances = dict(nx.all_pairs_shortest_path_length(subgraph_V, cutoff=3))
        nodes_list = list(subgraph_V.nodes)
        num_nodes = len(subgraph_V)
        A_V = np.zeros((num_nodes, num_nodes))
        for i in range(num_nodes):
            id1 = nodes_list[i]
            if id1 in distances:
                for j in range(num_nodes):
                    id2 = nodes_list[j]
                    if id2 in distances[id1] and id1 != id2:
                        A_V[i, j] = 1
        distances = dict(nx.all_pairs_shortest_path_length(subgraph_H, cutoff=3))
        nodes_list = list(subgraph_H.nodes)
        num_nodes = len(subgraph_H)
        A_H = np.zeros((num_nodes, num_nodes))
        for i in range(num_nodes):
            id1 = nodes_list[i]
            if id1 in distances:
                for j in range(num_nodes):
                    id2 = nodes_list[j]
                    if id2 in distances[id1] and id1 != id2:
                        A_H[i, j] = 1

        loops = np.argwhere((A_H.dot(A_V) + A_V.dot(A_H) == 2) |
                            (A_H.transpose().dot(A_V) + A_V.dot(A_H.transpose()) == 2))
        loop_vertices = list(set(loops.flatten()))
        for x in loops:
            edges_from_node = [e for e in graph.out_edges(x[0], keys=True, data=True)
                               if e[3]["direction"] == "l"]
            lengthx = min(abs(all_nodes[x[0]]["x_0"]-all_nodes[x[1]]["x_1"]), abs(all_nodes[x[0]]["x_1"]-all_nodes[x[1]]["x_0"]))
            lengthy = min(abs(all_nodes[x[0]]["y_0"] - all_nodes[x[1]]["y_1"]), abs(all_nodes[x[0]]["y_1"] - all_nodes[x[1]]["y_0"]))
            length = np.sqrt(lengthx**2 + lengthy**2)
            if len(edges_from_node) == 0:
                graph.add_edge(x[0], x[1], length=length,
                               lengthx_phys=lengthx, lengthy_phys=lengthy, direction="l",
                               weight=1 / length, key="l")
            elif len(edges_from_node) > 0:
                if edges_from_node[0][3]["length"] >= length:
                    graph.remove_edges_from(edges_from_node)
                    graph.add_edge(x[0], x[1], length=length,
                                   lengthx_phys=lengthx, lengthy_phys=lengthy, direction="l",
                                   weight=1 / length, key="l")
        loop_labels.update(dict.fromkeys(loop_vertices, 1))
        nx.set_node_attributes(graph, loop_labels, 'is_loop')
        return graph

    def fulfills_horizontal_edge_conditions(self, node1, node2, page,
                                            delta_f, length_x_min, delta_y_min):
        """

        :param node1:
        :param node2:
        :param page:
        :param delta_f:
        :param length_x_min:
        :param delta_y_min:
        :return:
        """
        return (delta_y_min < self.meta.y_eps or
                (constants.USE_VERTICAL_OVERLAP and
                 GraphConverter.is_vertically_overlapping(node1, node2))) and (
                (length_x_min < self.meta.threshold_x[page] and
                 (not constants.USE_FONT or
                  delta_f < self.meta.font_eps_h
                  )) or
                (constants.USE_RECT and
                 RectangleUtil.is_same_rectangle(node1, node2)))

    def fulfills_vertical_edge_conditions(self, node1, node2, page,
                                          delta_x_min, length_y_min, delta_f):
        """

        :param node1:
        :param node2:
        :param page:
        :param delta_x_min:
        :param length_y_min:
        :param delta_f:
        :return:
        """
        width_1 = node1["x_1"] - node1["x_0"]
        width_2 = node2["x_1"] - node2["x_0"]
        width_rel = abs(width_1 - width_2) / max(width_1, width_2, 1)
        return (delta_x_min < self.meta.x_eps or
                (constants.USE_HORIZONTAL_OVERLAP and
                 GraphConverter.is_horizontally_overlapping(node1, node2))) and\
                (length_y_min < self.meta.threshold_y[page] and
                 (not constants.USE_WIDTH or
                  ((width_rel < self.meta.width_pct_eps or
                    max(width_1, width_2) < self.meta.threshold_page_width[page])))
                 and (delta_f < self.meta.font_eps_v or not constants.USE_FONT)
                 or (constants.USE_RECT and
                     RectangleUtil.is_same_rectangle(node1, node2)))

    @staticmethod
    def is_horizontally_overlapping(node1, node2):
        """

        :param node1:
        :param node2:
        :return:
        """
        union = max(node2["x_1"]-node1["x_0"], node1["x_1"]-node2["x_0"])
        intersection = min(node1["x_1"], node2["x_1"]) - max(node1["x_0"], node2["x_0"])
        return intersection / union >= constants.HORIZONTAL_OVERLAP_THRESHOLD

    @staticmethod
    def is_vertically_overlapping(node1, node2):
        """

        :param node1:
        :param node2:
        :return:
        """
        min_height = min(node1["y_0"] - node1["y_1"], node2["y_0"] - node2["y_1"])
        overlap = min(node1["y_0"], node2["y_0"]) - max(node1["y_1"], node2["y_1"])
        return overlap / min_height >= constants.VERTICAL_OVERLAP_THRESHOLD

    def get_media_boxes(self):
        """

        :return:
        """
        return self.media_boxes

    def get_page_count(self):
        """

        :return:
        """
        return self.n

    def get_meta(self):
        """

        :return:
        """
        if self.meta is None:
            self.convert()
        return self.meta

    def get_graphs(self):
        """

        :return:
        """
        if self.graphs is None:
            self.convert()
        return self.graphs

max_stars_repo_name: MBAigner/GraphConverter

max_issues_repo_name: MBAigner/GraphConverter

max_forks_repo_name: MBAigner/GraphConverter

reason: import numpy,import networkx


text:
import numpy as np
from source_ddc.simulation_tools import simulate
from source_ddc.algorithms import NFXP, NPL, CCP
from source_ddc.probability_tools import StateManager, random_ccp
from test.utils.functional_tools import average_out

n_repetitions = 10

def test_nfxp(simple_transition_matrix):

    def utility_fn(theta, choices, states):
        m_states, m_actions = np.meshgrid(states, choices)
        return (theta[0] * np.log(m_states + 1) - theta[1] * m_actions).reshape((len(choices), -1, 1))

    true_params = [0.5, 3]
    discount_factor = 0.95
    n_choices = 2
    n_states = 5

    state_manager = StateManager(miles=n_states)

    @average_out(n_repetitions)
    def test():
        df, _ = simulate(
            500,
            100,
            n_choices,
            state_manager,
            true_params,
            utility_fn,
            discount_factor,
            simple_transition_matrix
        )

        algorithm = NFXP(
            df['action'].values,
            df['state'].values,
            simple_transition_matrix,
            utility_fn,
            discount_factor,
            parameter_names=['variable_cost', 'replacement_cost']
        )

        return algorithm.estimate(start_params=[-1, -1], method='bfgs')

    mean_params = test()
    tolerance_levels = np.array([0.05, 0.05])
    assert np.all(np.abs(mean_params - true_params) < tolerance_levels)

def test_ccp(simple_transition_matrix):

    def utility_fn(theta, choices, states):
        m_states, m_actions = np.meshgrid(states, choices)
        return (theta[0] * np.log(m_states + 1) - theta[1] * m_actions).reshape((len(choices), -1, 1))

    true_params = [0.5, 3]
    discount_factor = 0.95
    n_choices = 2
    n_states = 5

    state_manager = StateManager(miles=n_states)

    @average_out(n_repetitions)
    def test():
        df, ccp = simulate(
            500,
            100,
            n_choices,
            state_manager,
            true_params,
            utility_fn,
            discount_factor,
            simple_transition_matrix
        )

        algorithm = CCP(
            df['action'].values,
            df['state'].values,
            simple_transition_matrix,
            utility_fn,
            discount_factor,
            initial_p=ccp,
            parameter_names=['variable_cost', 'replacement_cost']
        )

        return algorithm.estimate(start_params=[1, 1], method='bfgs')

    mean_params = test()
    tolerance_levels = np.array([0.05, 0.05])
    assert np.all(np.abs(mean_params - true_params) < tolerance_levels)

def test_npl(simple_transition_matrix):

    def utility_fn(theta, choices, states):
        m_states, m_actions = np.meshgrid(states, choices)
        return (theta[0] * np.log(m_states + 1) - theta[1] * m_actions).reshape((len(choices), -1, 1))

    true_params = [0.5, 3]
    discount_factor = 0.95
    n_choices = 2
    n_states = 5

    state_manager = StateManager(miles=n_states)

    @average_out(n_repetitions)
    def test():
        df, _ = simulate(
            500,
            100,
            n_choices,
            state_manager,
            true_params,
            utility_fn,
            discount_factor,
            simple_transition_matrix)

        ccp = random_ccp(n_states, n_choices)

        algorithm = NPL(
            df['action'].values,
            df['state'].values,
            simple_transition_matrix,
            utility_fn,
            discount_factor,
            initial_p=ccp,
            parameter_names=['variable_cost', 'replacement_cost']
        )
        return algorithm.estimate(start_params=[1, 1], method='bfgs')

    mean_params = test()
    tolerance_levels = np.array([0.05, 0.05])
    assert np.all(np.abs(mean_params - true_params) < tolerance_levels)

def test_npl_relaxation_param(simple_transition_matrix):

    def utility_fn(theta, choices, states):
        m_states, m_actions = np.meshgrid(states, choices)
        return (theta[0] * np.log(m_states + 1) - theta[1] * m_actions).reshape((len(choices), -1, 1))

    true_params = [0.5, 3]
    discount_factor = 0.95
    n_choices = 2
    n_states = 5

    state_manager = StateManager(miles=n_states)

    @average_out(n_repetitions)
    def test():
        df, _ = simulate(500,
                         100,
                         n_choices,
                         state_manager,
                         true_params,
                         utility_fn,
                         discount_factor,
                         simple_transition_matrix)

        ccp = random_ccp(n_states, n_choices)

        algorithm = NPL(
            df['action'].values,
            df['state'].values,
            simple_transition_matrix,
            utility_fn,
            discount_factor,
            initial_p=ccp,
            relaxation_param=0.9,
            parameter_names=['variable_cost', 'replacement_cost'],
            npl_maxiter=50
        )

        return algorithm.estimate(start_params=[1, 1], method='bfgs')

    mean_params = test()
    tolerance_levels = np.array([0.05, 0.05])
    assert np.all(np.abs(mean_params - true_params) < tolerance_levels)

max_stars_repo_name: sansan-inc/econ-source

max_issues_repo_name: sansan-inc/econ-source

max_forks_repo_name: sansan-inc/econ-source


text:
import numpy as np
from tqdm import tqdm
from torch.utils.data import Dataset
import torch
import torchvision.transforms as transforms
from torch import Tensor
import csv
import os
from PIL import Image

class OpenSurfaces(Dataset):
  def __init__(self, root_dir, transform=None):
    self.root_dir = root_dir
    self.transform = transform
    self.mean = Tensor([0.48485812, 0.40377784, 0.32280155])
    self.std = Tensor([0.37216536, 0.349832,   0.37452201])
    self.classes = ["Wood", "Concrete", "Granite/marble", "Tile"]
    self.name_to_class = {c:idx for idx, c in enumerate(self.classes)}
    shape_indexes = {idx:[] for idx, c in enumerate(self.classes)}
    file_name = os.path.join(root_dir, 'shapes.csv')

    self.data = []
    with open(file_name,'r') as f:
      csv_reader = csv.reader(f)
      next(csv_reader)
      shape_idx_offset = 0
      for shape_idx, shape in enumerate(csv_reader):
        if shape[2] in self.classes:
          image_path = os.path.join("shapes-cropped", shape[0] + ".png")
          shape_class = self.name_to_class[shape[2]]
          shape_indexes[shape_class].append(shape_idx - shape_idx_offset)
          self.data.append([image_path, shape_class])
    l = len(self.data)

  def __len__(self):
    return len(self.data)

  def __getitem__(self, idx):
    image_path, shape_class = self.data[idx]
    total_path = os.path.join(self.root_dir, image_path)
    if not os.path.exists(total_path):
      print("No exist: " + image_path)
      return None, None
    image = Image.open(total_path)
    image = image.convert('RGBA')
    if self.transform:
      image = self.transform(image)
      #if is_tensor(image):
      #  image = transforms.Normalize(self.mean, self.std)(image)
    #print(image_path)
    return image, shape_class

class CropSampler(torch.nn.Module):
  def __init__(self, size):
    super().__init__()
    self.size = size

  def forward(self, img):
      width, height = img.size
      if width <  self.size or height <  self.size:
        return -1
      np_image = np.array(img)
      content_mask = (np_image != [255,255,255,0])[:,:,0]
      sample_count_width = int(width/self.size)
      sample_count_height = int(height/self.size)
      #half_size = int(self.size / 2)
      crop_area = self.size*self.size
      crops = []
      for i in range(sample_count_width):
        for j in range(sample_count_height):
          count = np.sum(content_mask[i*self.size : (i + 1)*self.size, j*self.size : (j + 1)*self.size])
          if count == crop_area:
            crops.append((np_image[i * self.size:(i + 1) * self.size, j * self.size:(j + 1) * self.size])[:,:,:3])
      if len(crops) == 0:
        return -1
      return np.array(crops)

sample_dataset = OpenSurfaces(root_dir="../datasets", transform=CropSampler(224))
sample_loader = torch.utils.data.DataLoader(dataset=sample_dataset, batch_size=1)
dataset_path = "../datasets/sample_224"
class_counter = {}
for c in sample_dataset.classes:
  class_counter[c] = 0

for images, label in tqdm(sample_loader):
  if len(images.shape) > 1:
    class_name = sample_dataset.classes[label]
    class_count = class_counter[class_name]
    if class_name == "Granite/marble":
      class_name = "Granite_marble"
    path_folder = os.path.join(dataset_path, class_name)
    for image in images[0]:
       img = Image.fromarray(np.array(image))
       path_image = os.path.join(path_folder, str(class_count) + ".png")
       img.save(path_image,"png")
       class_count += 1
    class_counter[sample_dataset.classes[label]] = class_count