text:
'''
simulate_PV_time_series: Model for generating PV soiling time series described in Deceglie et. al.
    "Numerical Validation of an Algorithm for Combined Soiling and Degradation Analysis of Photovoltaic
    Systems", Proceedings of the 2019 IEEE PVSC.

simulate_PV_time_series_gaussian_soiling: Model for generating PV soiling time series with soiling
    rates that vary from day to day according to a Gaussian distribution (Åsmund)

generate_my_soiling_signals: A function that generates 6 different time series based on the previous
    two functions. This was used to generate the datasets used in
    A. Skomedal, M. G. Deceglie, 2020, "Combined Estimation of Degradation and Soiling Losses in PV
    Systems"

Authors: Michael Deceglie and Åsmund Skomedal (Sep 2019)
'''

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def simulate_PV_time_series(first_date, last_date, freq='1D', degradation_rate=-0.005,
                            noise_scale=.01, seasonality_scale=.01, nr_of_cleaning_events=120,
                            soiling_rate_low=0.0001, soiling_rate_high=0.003,
                            smooth_rates=False, random_seed=False):

    if random_seed: # Have seed for repeatability
        if not type(np.random.seed) == int:
            np.random.seed(int(random_seed))

    # Initialize time series and data frame
    times = pd.date_range(first_date, last_date, freq=freq)
    df = pd.DataFrame(index=times, columns=['day', 'noise', 'seasonality', 'degradation', 'soiling',
                                            'cleaning_events'])
    n = len(times)

    df['day'] = range(n)
    df['noise'] = np.random.normal(1.0, scale=noise_scale, size=n)
    df['seasonality'] = seasonality_scale * np.sin(df['day'] / 365.25 * 2 * np.pi) + 1
    df['degradation'] = 1 + degradation_rate / 365.25 * df['day']

    # simulate soiling
    cleaning_events = np.random.choice(df['day'], nr_of_cleaning_events, replace='False')
    cleaning_events = np.sort(cleaning_events)

    x = np.full(n, 1)
    intervals = np.split(x, cleaning_events)
    soiling_rate = []
    for interval in intervals:
        rate = np.random.uniform(low=soiling_rate_low, high=soiling_rate_high, size=1) * np.ones(len(interval))
        soiling_rate.append(rate)
    df['soiling_rate'] = np.concatenate(soiling_rate)
    if smooth_rates:
        df.soiling_rate = df.soiling_rate.rolling(smooth_rates,center=True).mean().ffill().bfill()
    derate = np.concatenate([np.cumsum(si) for si in np.split(df['soiling_rate'].values, cleaning_events)])
    df['soiling'] = 1 - derate

    # Generate Performance Indexes
    df['PI_no_noise'] = df['seasonality'] * df['degradation'] * df['soiling']
    df['PI_no_soil'] = df['seasonality'] * df['degradation'] * df['noise']
    df['PI_no_degrad'] = df['seasonality'] * df['soiling']
    df['daily_norm'] = df['noise'] * df['PI_no_noise']

    return df

def simulate_PV_time_series_seasonal_soiling(first_date, last_date, freq='1D',
                                             degradation_rate=-0.005, noise_scale=.01,
                                             seasonality_scale=.01, nr_of_cleaning_events=120,
                                             soiling_rate_center=0.002, soiling_rate_std=0.001,
                                             soiling_seasonality_scale=.9, random_seed=False,
                                             smooth_rates=False, seasonal_rates=False):
    ''' As the name implies, this function models soiling rates that vary from day to day according
    to a Gaussian distribution '''
    if random_seed: # Have seed for repeatability
        if not type(np.random.seed) == int:
            np.random.seed(int(random_seed))

    # Initialize time series and data frame
    times = pd.date_range(first_date, last_date, freq=freq)
    df = pd.DataFrame(index=times)
    n = len(times)

    df['day'] = range(n)
    df['noise'] = np.random.normal(1.0, scale=noise_scale, size=n)
    df['seasonality'] = seasonality_scale * np.sin(df['day'] / 365.25 * 2 * np.pi) + 1
    df['degradation'] = 1 + degradation_rate / 365.25 * df['day']

    soiling_seasonality = soiling_rate_center * soiling_seasonality_scale * \
                            np.sin(df['day']/365.25*2*np.pi + np.random.uniform(10, size=1)*np.pi)
    cleaning_probability = ((2-soiling_seasonality_scale)*soiling_seasonality.max()+soiling_seasonality)
    cleaning_probability /= cleaning_probability.sum()

    # simulate soiling
    cleaning_events = np.random.choice(df['day'].values, nr_of_cleaning_events, replace='False',
                                       p=cleaning_probability.values)
    cleaning_events = np.sort(cleaning_events)
    df['cleaning_events'] = [True if day in cleaning_events else False for day in df.day]

    x = np.full(n, 1)
    intervals = np.split(x, cleaning_events)
    soiling = []
    soiling_rate = []
    for interval in intervals:
        if len(soiling)>0:
            pos = len(np.concatenate(soiling))
        else:
            pos = 0
        if seasonal_rates:
            if smooth_rates:
                soiling_season = soiling_seasonality[pos:pos+len(interval)]
                rate = np.random.normal(soiling_rate_center+soiling_season, soiling_rate_std,
                                        size=len(interval))
            else:
                soiling_season = soiling_seasonality[pos]
                rate = np.random.normal(soiling_rate_center+soiling_season, soiling_rate_std,
                                        size=1) * np.ones(len(interval))
        else:
            if smooth_rates: print('Smooth rates not possible without seasonal rates')
            soiling_season = 0
            rate = np.random.normal(soiling_rate_center+soiling_season, soiling_rate_std, size=1) \
                                    * np.ones(len(interval))
        derate = 1 - np.matmul(half_one_matrix(len(interval)), rate)
        soiling.append(derate)
        soiling_rate.append(rate)
    df['soiling'] = np.concatenate(soiling)
    df['soiling_rate'] = np.concatenate(soiling_rate)

    # Generate Performance Indexes
    df['PI_no_noise'] = df['seasonality'] * df['degradation'] * df['soiling']
    df['PI_no_soil'] = df['seasonality'] * df['degradation'] * df['noise']
    df['PI_no_degrad'] = df['seasonality'] * df['soiling']
    df['daily_norm'] = df['noise'] * df['PI_no_noise']

    return df

def generate_my_soiling_signals(high=.02, low=.01, index_form='Datetime'):
    first_date='2010/01/01'
    last_date='2020/01/01'
    noises = [low, low, high, low, low, low]
    seasonality_scales = [low, high, low, low, low, low]
    soiling_levels = [.003, .001, .001, .0005, .0005, .0001]
    names = ['normal', 'M Soil, H season', 'M Soil, H noise', 'Seasonal cleaning',
             'L/M Soil (.0005)', 'Low Soil (.0001)']
    dfs = []
    for i in range(6):
        if i==3:
            df = simulate_PV_time_series_seasonal_soiling(first_date, last_date,
                                                          noise_scale=noises[i],
                                                          seasonality_scale=seasonality_scales[i],
                                                          soiling_rate_std=0,
                                                          soiling_rate_center=soiling_levels[i],
                                                          soiling_seasonality_scale=.95,
                                                          smooth_rates=False,
                                                          seasonal_rates=False)
        else:
            df = simulate_PV_time_series(first_date, last_date,
                                         noise_scale=noises[i],
                                         seasonality_scale=seasonality_scales[i],
                                         nr_of_cleaning_events=120,
                                         soiling_rate_low=0,
                                         soiling_rate_high=soiling_levels[i],
                                         degradation_rate=-.005)
        if index_form=='numeric':
            df.index = np.arange(len(df))

        dfs.append(df)

    return dfs, names

def half_one_matrix(size):
    dummy = np.zeros((size,size))
    dummy[np.triu_indices(size, k=1)]=1
    return dummy.T

if __name__=='__main__':
    # Visualize my soiling signals
    np.random.seed(21)
    dfs, names = generate_my_soiling_signals()
    names = ['Base case \n(max soiling rate = 0.3 %/d)',
             'Double seasonality \n(max soiling rate = 0.1 %/d)',
             'Double noise \n(max soiling rate = 0.1 %/d)',
             'Seasonal cleaning \n(constant soiling rate = 0.05 %/d)',
             'Little soiling \n(max soiling rate = 0.05 %/d)',
             'Very little soiling \n(max soiling rate = 0.01 %/d)']
    N = len(names)
    spn = 5 if N in [5, 9, 10] else (3 if N in [6] else 4)
    plt.close('all')
    for i in range(N):
        df = dfs[i]
        if i%spn==0: fig, ax = plt.subplots(spn, 1, figsize=(10,6), sharex=True)
        ax[i%spn].scatter(df.index, df.daily_norm, 5, alpha=.5)
        ax[i%spn].plot(df.index, df.PI_no_noise, 'k', alpha=.5)
        ax[i%spn].set_ylim(.8,1.06)
        ax[i%spn].set_ylabel('PI')
        ax[i%spn].set_title(names[i])
        plt.tight_layout()


text:
"""This module defines utility functions for drawing."""
# Note for developers: this module should be imported only a specific function is needed.
# Please do not import this module when initializing ProDy for robustness reasons.

import numpy as np

from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d

__all__ = ['Arrow3D', 'drawArrow3D']

class Arrow3D(FancyArrowPatch):
    """This function is implemented by tacaswell on stackoverflow:
    https://stackoverflow.com/a/29188796."""
    def __init__(self, xs, ys, zs, *args, **kwargs):
        super(Arrow3D, self).__init__((0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def draw(self, renderer):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
        FancyArrowPatch.draw(self, renderer)

def drawArrow3D(*xyz, **kwargs):
    from matplotlib.pyplot import gca

    allargs = xyz
    xs = []; ys = []; zs = []; args = []

    for arg in allargs:
        if not np.isscalar(arg):
            if len(arg) == 3:
                xs.append(arg[0])
                ys.append(arg[1])
                zs.append(arg[2])
                continue
        args.append(arg)

    ax = gca()
    if 'arrowstyle' not in kwargs:
        kwargs['arrowstyle'] = '-|>'
    if 'mutation_scale' not in kwargs:
        kwargs['mutation_scale'] = 20
    if 'color' not in kwargs:
        kwargs['color'] = 'k'

    arrow = Arrow3D(xs, ys, zs, *args, **kwargs)
    ax.add_artist(arrow)
    return arrow

ext: py

max_stars_repo_name: fongchun/ProDy

max_issues_repo_name: fongchun/ProDy

max_forks_repo_name: fongchun/ProDy

reason: import numpy


text:
# import libraries
import sys

import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier, GradientBoostingClassifier, ExtraTreesClassifier
from sklearn.svm import SVC
from sklearn.multioutput import MultiOutputClassifier
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.model_selection import train_test_split
import sklearn.metrics as metrics
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report
from sklearn.linear_model import SGDClassifier, LogisticRegression
from sklearn.model_selection import GridSearchCV
from sklearn.decomposition import TruncatedSVD
from sklearn.base import BaseEstimator
import pickle
import re
import matplotlib.pyplot as plt
import seaborn as sns

import nltk
nltk.download(['punkt', 'wordnet'])

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

#imblearn
from imblearn.over_sampling import RandomOverSampler
from sklearn.base import BaseEstimator, TransformerMixin
from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit
from imblearn.under_sampling import RandomUnderSampler

def load_data(database_filepath):
    """
    Load the data

    Inputs:
    database_filepath: String. Filepath for the db file containing the cleaned data.

    Output:
    X: dataframe. Contains the feature data.
    y: dataframe. Contains the labels (categories) data.
    category_names: List of strings. Contains the labels names.
    """

    # load data from database
    engine = create_engine('sqlite:///{}'.format(database_filepath))
    df = pd.read_sql_table('messages', con=engine)

    X = df['message']
    y = df.drop(['message', 'genre', 'id', 'original'], axis=1)
    category_names = y.columns.tolist()

    return X, y, category_names

def tokenize(text):
    """
    Normalize, tokenize and stems texts.

    Input:
    text: string. Sentence containing a message.

    Output:
    stemmed_tokens: list of strings. A list of strings containing normalized and stemmed tokens.
    """
    url_regex = 'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'

    # get list of all urls using regex
    detected_urls = re.findall(url_regex, text)

    # replace each url in text string with placeholder
    for url in detected_urls:
        text = text.replace(url, "urlplaceholder")

    # tokenize text
    tokens = word_tokenize(text)

    # initiate lemmatizer
    lemmatizer = WordNetLemmatizer()

    # iterate through each token
    clean_tokens = []
    for tok in tokens:

        # lemmatize, normalize case, and remove leading/trailing white space
        clean_tok = lemmatizer.lemmatize(tok).lower().strip()
        clean_tokens.append(clean_tok)

    return clean_tokens

class StartVerbExtractor(BaseEstimator, TransformerMixin):
    """
    Creates feature based on verb presence in text
    """
    def start_verb(self, text):
        sentence_list = nltk.sent_tokenize(text)
        for sentence in sentence_list:
            pos_tags = nltk.pos_tag(tokenize(sentence))
            if len(pos_tags) != 0:
                first_word, first_tag = pos_tags[0]
                if first_tag in ['VB', 'VBP'] or first_word == 'RT':
                    return 1
        return 0

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X_tag = pd.Series(X).apply(self.start_verb)
        return pd.DataFrame(X_tag)

def get_text_len(data):
    """
    Gets the text length.
    """

    return np.array([len(text) for text in data]).reshape(-1, 1)

def build_model():
    """
    Builds a ML pipeline and performs gridsearch.
    Args:
    None
    Returns:
    cv: gridsearchcv object.
    """

    # Creates pipeline
    pipeline = Pipeline([
    ('features', FeatureUnion([
            ('text_pipeline', Pipeline([
                ('vect', CountVectorizer(tokenizer=tokenize)),
                ('best', TruncatedSVD()),
                ('tfidf', TfidfTransformer())])),
        ('start_verb', StartVerbExtractor())])),
    ('clf', MultiOutputClassifier(AdaBoostClassifier(random_state=42)))
    ])

    # Parameters
    parameters = {'clf__estimator__n_estimators': [100, 200, 300],
              'clf__estimator__random_state': [42],
             'clf__estimator__learning_rate': [0.05]}

    # Creates gridsearch
    cv = GridSearchCV(pipeline, param_grid = parameters, cv = 10, refit = True, verbose = 1, return_train_score = True, n_jobs = -1)

    return cv

def evaluate_model(model, X_test, y_test, category_names):
    """
    Returns test accuracy, number of 1s and 0s, recall, precision and F1 Score.

    Inputs:
    model: model object. Instanciated model.
    X_test: pandas dataframe containing test features.
    y_test: pandas dataframe containing test labels.
    category_names: list of strings containing category names.

    Returns:
    None
    """

    y_pred = model.predict(X_test)
    y_pred_pd = pd.DataFrame(y_pred, columns = category_names)

    print(classification_report(y_test, y_pred, target_names = category_names))

def save_model(model, model_filepath):

    pickle.dump(model.best_estimator_, open(model_filepath, 'wb'))

def main():
    if len(sys.argv) == 3:
        database_filepath, model_filepath = sys.argv[1:]
        print('Loading data...\n    DATABASE: {}'.format(database_filepath))
        X, y, category_names = load_data(database_filepath)

        mlss = MultilabelStratifiedShuffleSplit(n_splits=1, test_size=0.33, random_state=42)

        for train_index, test_index in mlss.split(X, y):
            X_train, X_test = X[train_index], X[test_index]
            y_train, y_test = y.values[train_index], y.values[test_index]
        y_train = pd.DataFrame(y_train,columns=category_names)
        y_test = pd.DataFrame(y_test,columns=category_names)

        print('Building model...')
        model = build_model()

        print('Training model...')
        model.fit(X_train, y_train)

        print('Evaluating model...')
        evaluate_model(model, X_test, y_test, category_names)

        print('Saving model...\n    MODEL: {}'.format(model_filepath))
        save_model(model, model_filepath)

        print('Trained model saved!')

    else:
        print('Please provide the filepath of the disaster messages database '\
              'as the first argument and the filepath of the pickle file to '\
              'save the model to as the second argument. \n\nExample: python '\
              'train_classifier.py ../data/DisasterResponse.db classifier.pkl')

if __name__ == '__main__':
    main()

max_stars_repo_name: sousablde/Disaster-Response-Pipeline

max_stars_repo_licenses: CNRI-Python blessing

max_issues_repo_name: sousablde/Disaster-Response-Pipeline

max_issues_repo_licenses: CNRI-Python blessing

max_forks_repo_name: sousablde/Disaster-Response-Pipeline

max_forks_repo_licenses: CNRI-Python blessing


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

from numpy import *

class solution_zero_one_bag_weight:

    def run_recursive_cache(self,weights,capacity):
        """
        回溯法 解 01背包问题 , 并缓存子问题的解
        ref: https://time.geekbang.org/column/article/74788
        :param weights:
        :param capacity:
        :return:
        """
        self.weights=weights
        self.capacity=capacity

        self.max_bag_weight=0
        # self.res_bag=()

        self.cache = zeros((len(weights)+1 , capacity+1), dtype=bool) # 默认为 False

        i=0
        current_bag_weight=0
        self.__f( i, current_bag_weight)

        return self.max_bag_weight

    def __f(self, i , current_bag_weight):

        if i <= len(self.weights): # i= 0,1,2,3

            if current_bag_weight > self.capacity: #搜索剪枝: 当发现已经选择的物品的重量超过 Wkg 之后，我们就停止继续探测剩下的物品
                return

            if self.cache[i][current_bag_weight] == True: # 判断值相等 用 '=='
                print('(',i,',',current_bag_weight,')','has cached')
                return
            else:
                self.cache[i][current_bag_weight] = True

            print(i, current_bag_weight) # 树的先序列遍历

            if current_bag_weight > self.max_bag_weight:
                self.max_bag_weight =current_bag_weight

            self.__f(i+1,current_bag_weight) # f(1,0)

            if i < len(self.weights):
                self.__f(i+1, current_bag_weight + self.weights[i] ) #f(1,2)

    def run_dynamic_programming(self,weights,capacity):
        """
        动态规划法 解 01背包问题
        ref: https://time.geekbang.org/column/article/74788
        :param weights:
        :param capacity:
        :return:
        """
        states = zeros((len(weights), capacity + 1), dtype=int)

        row_num=shape(states)[0]
        col_num=shape(states)[1]

        for i in range(row_num):

            states[i][0]=1 # 第i个物品 不放入背包

            if  weights[i] < capacity :
                states[i][weights[i]]=1 # 第i个物品放入背包

            for j in range(col_num):
                if states[i-1][j]==1:
                    states[i][j] = 1 # 第i个物品 不放入背包

                    if j+weights[i] < capacity + 1:
                        states[i][j+weights[i]] = 1 # 第i个物品放入背包

        print(states)
        res_index=0
        for i  in  range(col_num-1,-1,-1): # 从后往前遍历列，找到第一个为1的，即为 最大物品重量和
            if states[-1][i]==1:
                res_index=i
                break

        return res_index

    def run_dynamic_programming_v2(self,weights,capacity):
        """
        动态规划法v2 解 01背包问题：
        （1）降低 空间复杂度， O(n*w) -> O(2w)

        ref: https://time.geekbang.org/column/article/74788
        :param weights:
        :param capacity:
        :return:
        """
        states = zeros((capacity + 1), dtype=int)

        row_num=len(weights)
        col_num=capacity + 1

        for i in range(row_num):

            next_states = zeros((capacity + 1), dtype=int)
            next_states[0]=1 # 第i个物品 不放入背包

            if  weights[i] < capacity :
                next_states[weights[i]]=1 # 背包中 只有 第i个物品

            for j in range(col_num):
                if states[j]==1:
                    next_states[j] = 1 # 第i个物品 不放入背包

                    if j+weights[i] < capacity + 1:
                        next_states[j+weights[i]] = 1 # 第i个物品放入背包

            states=next_states

        print(states)
        res_index=0
        for i  in  range(col_num-1,-1,-1):
            if states[i]==1:
                res_index=i
                break

        return res_index

    def run_dynamic_programming_v3(self, weights, capacity):
        """
        动态规划法v3 解 01背包问题：
        （1） 进一步降低 空间复杂度: O(2w) -> O(w)，
         (2) 降低时间复杂度: 对 states 的列的遍历 做 剪枝
         (3) 当重复使用一个 states 时，对 states 的列的遍历 j 需要从大到小来处理。如果我们按照 j 从小到大处理的话，会出现 for 循环重复计算的问题。

        ref: https://time.geekbang.org/column/article/74788
        :param weights:
        :param capacity:
        :return:
        """
        states = zeros((capacity + 1), dtype=int)

        row_num = len(weights)
        col_num = capacity + 1

        # 对状态0 特殊处理：
        states[0] = 1
        if weights[0] < capacity + 1:
            states[weights[0]] = 1  # 第0个物品放入背包

        for i in range(1,row_num): # 从状态1开始
            col_end=capacity-(weights[i]) # 剪枝: 大于col_end 的j 不用考虑，因为 j + weights[i] 必然大于 背包的容量
            for j in range(col_end,-1,-1):
                if states[j] == 1:
                    states[j + weights[i]] = 1  # 第i个物品放入背包

        print(states)

        res_index = 0
        for i in range(col_num - 1, -1, -1):
            if states[i] == 1:
                res_index = i
                break

        return res_index

import sys
class  solution_double11advance:

    def run_dynamic_programming_v2(self, prices, condition):
        """
        双十一 购物节有 “满 200 元减 50 元 ” 的促销。假设女朋友的购物车中有 n 个（n>100）想买的商品，她希望从里面选几个，
        在凑够满减条件的前提下，让选出来的商品价格总和最大程度地接近满减条件（200 元），这样就可以极大限度地“薅羊毛”。
        请给出选择商品的方案。

        空间复杂度： O(2*w) w为 conditionx2

        :param prices: 各个商品的价格
        :param condition: 满减的条件
        :return: 选择的商品方案
        """
        capacity = condition * 2

        states = zeros(capacity + 1, dtype=int)

        print('states memory_size:', sys.getsizeof(states))

        row_num = len(prices)
        col_num = condition + 1

        bag_states = []

        for i in range( row_num):  # 从第0个状态开始

            next_states = zeros((capacity + 1), dtype=int)

            next_states[0] = 1  # 第i个物品 不放入背包

            if prices[i] < capacity:
                next_states[prices[i]] = 1  # 背包中 只有 第i个物品

            for j in range(col_num):
                if states[j] == 1:
                    next_states[j] = 1  # 第i个物品 不放入背包

                    if j + prices[i] < capacity + 1:
                        next_states[j + prices[i]] = 1  # 第i个物品放入背包

            states = next_states

            bag_states.append([index for index in range(capacity + 1) if
                               states[index] != 0])  # 记录 各个状态i 下 背包的总价

        print('next_states memory_size:', sys.getsizeof(next_states))

        print(bag_states)  #
        print('bag_states memory_size:', sys.getsizeof(bag_states))

        min_price = 0
        for i in range(condition, capacity + 1):  # 满减问题中，我们要找 大于等于 200（满减条件）的值中最小的
            if states[i] == 1:
                min_price = i
                break

        # 根据背包中物品的总价，反推出 背包中的物品
        bag = []
        bag_price = min_price

        for i in range(row_num - 1, -1, -1):
            if i > 0:
                if bag_price - prices[i] in bag_states[i - 1] :
                    # case1: 第i个物品装入背包
                    # 1.bag_states[i-1] 上一个状态 背包的 总价的情况
                    # 3. (当前背包的重量 - 第i个物品的价格) 为 上一个状态的 背包的总价 则说明 第i 个物品被放入背包

                    bag.append(i)
                    bag_price = bag_price - prices[i]

                else:
                    # case2: 第i个物品未装入背包,背包的总价不变
                    pass

            elif i == 0:
                if bag_price == prices[i]:
                    bag.append(i)

        return min_price, bag

    def run_dynamic_programming_v3(self, prices, condition):
        """
        空间复杂度： O(w)
        :param prices: 各个商品的价格
        :param condition: 满减的条件
        :return: 选择的商品方案
        """
        capacity = condition * 2

        states = zeros(capacity + 1, dtype=int)

        print('states memory_size:', sys.getsizeof(states))

        row_num = len(prices)
        col_num = condition + 1

        bag_states = []

        states[0] = 1  # 第0个 状态 特殊处理
        if prices[0] < capacity + 1:
            states[prices[0]] = 1  # 第0个物品放入背包

        for i in range(1, row_num):  # 从第1个状态开始

            col_end = capacity - (prices[i])  # 剪枝
            for j in range(col_end, -1, -1):
                if states[j] == 1:
                    states[j + prices[i]] = 1  # 第i个物品放入背包

            bag_states.append([index for index in range(capacity + 1) if
                               states[index] != 0])  # 记录 各个状态i 下 背包的总价的所有情况

        print(bag_states)  #
        print('bag_states memory_size:', sys.getsizeof(bag_states))

        min_price = 0
        for i in range(condition, capacity + 1):  # 满减问题中，我们要找 大于等于 200（满减条件）的值中最小的
            if states[i] == 1:
                min_price = i
                break

        # 根据背包中物品的总价，反推出 背包中的物品
        bag = []
        bag_price = min_price

        for i in range(row_num - 1, -1, -1):
            if i > 0:
                if bag_price - prices[i] in bag_states[i - 1] :
                    # case1: 第i个物品装入背包
                    # 1.bag_states[i-1] 上一个状态 背包的 总价的情况
                    # 3. (当前背包的重量 - 第i个物品的价格) 为 上一个状态的 背包的总价 则说明 第i 个物品被放入背包

                    bag.append(i)
                    bag_price = bag_price - prices[i]

                else:
                    # case2: 第i个物品未装入背包,背包的总价不变
                    pass

            elif i == 0:
                if bag_price == prices[i]:
                    bag.append(i)

        return min_price, bag

    def run_dynamic_programming_v4(self, prices, condition):
        """
        存储 动态规划 状态的数组 过于稀疏，尝试用其他数据结构  dict 替代 数组，进一步降低内存占用
        空间复杂度： 与 condition 的大小无关

        对比 bag_states 的内存空间：
         v3 :states memory_size: 3300
         v4 :dict states memory_size: 1184

        :param prices: 各个商品的价格
        :param condition: 满减的条件
        :return: 选择的商品方案
        """
        capacity = condition * 2

        states = {}

        print('dict states(None) memory_size:', sys.getsizeof(states))

        print(' prices memory_size:', sys.getsizeof(prices))

        row_num = len(prices)
        col_num = condition + 1

        bag_states = []

        for i in range(row_num):  # 从第0个状态开始

            next_states = {}

            next_states[0] = 1  # 第i个物品 不放入背包

            if prices[i] < capacity:
                next_states[prices[i]] = 1  # 背包中 只有 第i个物品

            for j in range(col_num):
                if states.get(j)==1:
                    next_states[j] = 1  # 第i个物品 不放入背包

                    if j + prices[i] < capacity + 1:
                        next_states[j + prices[i]] = 1  # 第i个物品放入背包

            states = next_states

            bag_states.append(states.keys())  # 记录 各个状态i 下 背包的总价

        print('dict states memory_size:', sys.getsizeof(states))

        print(bag_states)  #
        print('bag_states memory_size:', sys.getsizeof(bag_states))

        min_price = 0
        for i in range(condition, capacity + 1):  # 满减问题中，我们要找 大于等于 200（满减条件）的值中最小的
            if states.get(i)==1:
                min_price = i
                break

        # 根据背包中物品的总价，反推出 背包中的物品
        bag = []
        bag_price = min_price

        for i in range(row_num - 1, -1, -1):
            if i > 0:
                if bag_price - prices[i] in bag_states[i - 1]:
                    # case1: 第i个物品装入背包
                    # 1.bag_states[i-1] 上一个状态 背包的 总价的情况
                    # 3. (当前背包的重量 - 第i个物品的价格) 为 上一个状态的 背包的总价 则说明 第i 个物品被放入背包

                    bag.append(i)
                    bag_price = bag_price - prices[i]

                else:
                    # case2: 第i个物品未装入背包,背包的总价不变
                    pass

            elif i == 0:
                if bag_price == prices[i]:
                    bag.append(i)

        return min_price, bag

class solution_zero_one_bag_value:

    def run_recursive_cache(self, weights, values, capacity):
        """
        回溯法（子问题缓存） 解  01背包问题 升级版（上价值）
        ref: https://time.geekbang.org/column/article/74788
        :param weights:
        :param values:
        :param capacity:
        :return:
        """
        self.weights = weights
        self.values = values

        self.capacity = capacity

        self.max_bag_value = 0

        self.cache_value = zeros((len(weights) + 1, capacity + 1), dtype=int)

        i = 0
        current_bag_weight = 0
        current_bag_value=0

        self.__f(i, current_bag_weight,current_bag_value)

        return self.max_bag_value

    def __f(self, i, current_bag_weight,current_bag_value):

        if i <= len(self.weights):  # i= 0,1,2,3

            if current_bag_weight > self.capacity:  # 搜索剪枝: 当发现已经选择的物品的重量超过 Wkg 之后，我们就停止继续探测剩下的物品
                return

            if current_bag_value<self.cache_value[i][current_bag_weight] :
                print('(', i, ',', current_bag_weight,',', ')', 'has cached',', the value is ',self.cache_value[i][current_bag_weight])
                return
            else:
                self.cache_value[i][current_bag_weight] = current_bag_value

            print(i, current_bag_weight,current_bag_value)  # 树的先序列遍历

            if current_bag_value > self.max_bag_value:
                self.max_bag_value = current_bag_value

            self.__f(i + 1, current_bag_weight,current_bag_value)  # f(1,0,0)

            if i < len(self.weights):
                self.__f(i + 1, current_bag_weight + self.weights[i],current_bag_value+self.values[i])  # f(1,2,3)

    def run_dynamic_programming_v2(self, weights,values, capacity):
        """
        动态规划法v2 解 01背包问题（升级版） ：
        空间复杂度 ： O(2w)
        ref: https://time.geekbang.org/column/article/74788
        :param weights:
        :param values:
        :param capacity:
        :return:
        """
        states_value = zeros((capacity + 1), dtype=int)

        row_num = len(weights)
        col_num = capacity + 1

        bag_states=[]

        for i in range(row_num):

            next_states_value = zeros((capacity + 1), dtype=int) #当前背包中的物品价值之和

            if weights[i] <= capacity: next_states_value[weights[i]]=values[i] # 背包中 只有 第i 个物品

            # col_end=capacity-(weights[i])

            for j in range(col_num):

                if  states_value[j] > 0: #上一次 背包中的物品价值之和

                    if next_states_value[j] < states_value[j]: #把每一层中 (i, j) 重复的状态合并，只记录 物品价值 最大的那个状态
                        next_states_value[j]=states_value[j]  # 价值大才更新

                    if j + weights[i]<= capacity:
                        next_states_value[j + weights[i]] =states_value[j]+ values[i]

            states_value = next_states_value

            print(i,states_value)

            bag_states.append( [[0,0]] + [ [index,states_value[index]]  for index in range( capacity + 1 ) if states_value[index]!=0 ]) # 状态转移矩阵过于稀疏，只保留有值的记录

        print(bag_states) #   [col0 , col1]=[8,17] ( col1 为背包的重量 col2 为包内物品的价值)
        # [
        #   [[0, 0], [2, 3]],
        #   [[0, 0], [2, 6], [4, 9]],
        #   [[0, 0], [2, 6], [4, 9], [6, 14], [8, 17]]
        # ]

        #背包 最重的 那个状态 不一定是 价值最大的
        max_value=max(states_value)
        max_value_index=argmax(states_value)

        # 根据最大的背包价值，反推出 背包中的物品
        bag=[]
        bag_weight=max_value_index # 包中的物品价值最大时，背包的重量
        bag_value=max_value

        for i in range(row_num-1,-1,-1):
            if i>0:
                if  (bag_weight-weights[i] in  array(bag_states[i-1])[:,0] ) and (bag_value-values[i] in  array(bag_states[i-1])[:,1] ):
                # case1: 第i个物品装入背包
                # 1.bag_states[i-1] 上一个状态
                # 2. array(bag_states[i-1])[:,0] 上一个状态的 背包的重量; array(bag_states[i-1])[:,1] 上一个状态的 背包的价值
                # 3. (当前背包的重量 - 第i个物品的重量) 为 上一个状态的 背包的重量 同时 (当前背包的价值 - 第i个物品的价值) 为 上一个状态的 背包的价值，则说明 第i 个物品被放入背包

                    bag.append(i)
                    bag_weight=bag_weight-weights[i]
                    bag_value=bag_value-values[i]

                else:
                # case2: 第i个物品未装入背包,背包的重量和价值不变
                    pass

            elif i==0:
                if bag_weight==weights[i] and bag_value==values[i]:
                    bag.append(i)

        return max_value,bag

if __name__ == '__main__':

    weights=[2,2,4,6,3]
    # weights = [2,2,4]
    sol = solution_zero_one_bag_weight()
    # print(sol.run_recursive_cache(weights,8))
    # print(sol.run_recursive_cache(weights, 21))

    # print(sol.run_dynamic_programming_v3(weights,9))

    weights = [2, 2, 4]
    values=[3,6,8]

    sol2=solution_zero_one_bag_value()
    # print(sol2.run_recursive_cache(weights,values,5))
    # print(sol2.run_dynamic_programming_v2(weights, values, 5))
    print(sol2.run_dynamic_programming_v2(weights, values, 9))

    weights = [2, 2, 4, 6, 3]
    values=[3,4,8,9,6]

    # print(sol2.run_recursive_cache(weights, values, 17))
    # print(sol2.run_dynamic_programming_v2(weights, values, 17))

    prices=[10,20,400,100,50,60]
    sol3=solution_double11advance()

    # print(sol3.run_dynamic_programming_v3(prices,200))
    # print(sol3.run_dynamic_programming_v4(prices, 200))
    # print(sol3.run_dynamic_programming_v3(prices,400))
    # print(sol3.run_dynamic_programming_v4(prices, 400))

max_stars_repo_name: Xinrihui/Data-Structure-and-Algrithms

max_issues_repo_name: Xinrihui/Data-Structure-and-Algrithms

max_forks_repo_name: Xinrihui/Data-Structure-and-Algrithms

reason: from numpy


text:
import cv2 as cv
import numpy as np
from flask import (
    Blueprint, request, abort, render_template
)

from src.backend import ALLOWED_EXTENSIONS
from src.ml.image_processor import Photomath

bp = Blueprint('service', __name__)

photomath = Photomath.create_vanilla_photomath()

@bp.route('/', methods=['GET', 'POST'])
def index():
    """
    Single page of Photomath web app. Renders simple UI on which user can
    upload a single valid image and request it's solving.

    :return: HTML page
    """
    if request.method == 'POST':
        image = request.files['image']

        if image.filename == '':
            abort(422, 'No selected file')
        if ('.' in image.filename and image.filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS) is False:
            abort(422, 'File format not supported')

        image = cv.imdecode(np.frombuffer(image.read(), np.uint8), cv.IMREAD_UNCHANGED)

        img = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
        expr, result = photomath.process_photo(img)
        return render_template('service.html', solution=result)

    return render_template('service.html')

max_stars_repo_name: Marko-Misko/photomath

max_issues_repo_name: Marko-Misko/photomath

max_forks_repo_name: Marko-Misko/photomath


text:
#!/usr/bin/env python

"""
Utilities for automatically detecting equilibrated region of molecular simulations.

John D. Chodera <john.chodera@choderalab.org>
Sloan Kettering Institute
Memorial Sloan Kettering Cancer Center

LICENSE

This is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as
published by the Free Software Foundation, either version 2.1
of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this software. If not, see <http://www.gnu.org/licenses/>.

"""

import numpy as np

def statisticalInefficiency_multiscale(A_n, B_n=None, fast=True, mintime=3):
    """
    Compute the (cross) statistical inefficiency of (two) timeseries using multiscale method from Chodera.

    Parameters
    ----------
    A_n : np.ndarray, float
        A_n[n] is nth value of timeseries A.  Length is deduced from vector.
    B_n : np.ndarray, float, optional, default=None
        B_n[n] is nth value of timeseries B.  Length is deduced from vector.
        If supplied, the cross-correlation of timeseries A and B will be estimated instead of the
        autocorrelation of timeseries A.
    fast : bool, optional, default=False
        f True, will use faster (but less accurate) method to estimate correlation
        time, described in Ref. [1] (default: False).
    mintime : int, optional, default=3
        minimum amount of correlation function to compute (default: 3)
        The algorithm terminates after computing the correlation time out to mintime when the
        correlation function first goes negative.  Note that this time may need to be increased
        if there is a strong initial negative peak in the correlation function.

    Returns
    -------
    g : np.ndarray,
        g is the estimated statistical inefficiency (equal to 1 + 2 tau, where tau is the correlation time).
        We enforce g >= 1.0.

    Notes
    -----
    The same timeseries can be used for both A_n and B_n to get the autocorrelation statistical inefficiency.
    The fast method described in Ref [1] is used to compute g.

    References
    ----------
    [1] J. D. Chodera, W. C. Swope, J. W. Pitera, C. Seok, and K. A. Dill. Use of the weighted
        histogram analysis method for the analysis of simulated and parallel tempering simulations.
        JCTC 3(1):26-41, 2007.

    Examples
    --------

    Compute statistical inefficiency of timeseries data with known correlation time.

    >>> from pymbar import testsystems
    >>> A_n = testsystems.correlated_timeseries_example(N=100000, tau=5.0)
    >>> g = statisticalInefficiency_multiscale(A_n, fast=True)

    """

    # Create numpy copies of input arguments.
    A_n = np.array(A_n)

    if B_n is not None:
        B_n = np.array(B_n)
    else:
        B_n = np.array(A_n)

    # Get the length of the timeseries.
    N = A_n.size

    # Be sure A_n and B_n have the same dimensions.
    if(A_n.shape != B_n.shape):
        raise Exception('A_n and B_n must have same dimensions.')

    # Initialize statistical inefficiency estimate with uncorrelated value.
    g = 1.0

    # Compute mean of each timeseries.
    mu_A = A_n.mean()
    mu_B = B_n.mean()

    # Make temporary copies of fluctuation from mean.
    dA_n = A_n.astype(np.float64) - mu_A
    dB_n = B_n.astype(np.float64) - mu_B

    # Compute estimator of covariance of (A,B) using estimator that will ensure C(0) = 1.
    sigma2_AB = (dA_n * dB_n).mean()  # standard estimator to ensure C(0) = 1

    # Trap the case where this covariance is zero, and we cannot proceed.
    if(sigma2_AB == 0):
        raise ParameterError('Sample covariance sigma_AB^2 = 0 -- cannot compute statistical inefficiency')

    # Accumulate the integrated correlation time by computing the normalized correlation time at
    # increasing values of t.  Stop accumulating if the correlation function goes negative, since
    # this is unlikely to occur unless the correlation function has decayed to the point where it
    # is dominated by noise and indistinguishable from zero.
    t = 1
    increment = 1
    while (t < N - 1):

        # compute normalized fluctuation correlation function at time t
        C = np.sum(dA_n[0:(N - t)] * dB_n[t:N] + dB_n[0:(N - t)] * dA_n[t:N]) / (2.0 * float(N - t) * sigma2_AB)
        # Terminate if the correlation function has crossed zero and we've computed the correlation
        # function at least out to 'mintime'.
        if (C <= 0.0) and (t > mintime):
            break

        # Accumulate contribution to the statistical inefficiency.
        g += 2.0 * C * (1.0 - float(t) / float(N)) * float(increment)

        # Increment t and the amount by which we increment t.
        t += increment

        # Increase the interval if "fast mode" is on.
        if fast:
            increment += 1

    # g must be at least unity
    if (g < 1.0):
        g = 1.0

    # Return the computed statistical inefficiency.
    return g

def statisticalInefficiency_geyer(A_n, method='con'):
    """Compute the statistical inefficiency of a timeseries using the methods of Geyer.

    Parameters
    ----------
    A_n : np.ndarray, float
        A_n[n] is nth value of timeseries A.  Length is deduced from vector.
    method : str, optional, default='con'
        The method to use; matches notation from `initseq` from `mcmc` R package by Geyer.
        'pos' : initial positive sequence (IPS) estimator
        'dec' : initial monotone sequence (IMS) estimator
        'con' : initial convex sequence (ICS) estimator

    Returns
    -------
    g : np.ndarray,
        g is the estimated statistical inefficiency (equal to 1 + 2 tau, where tau is the correlation time).
        We enforce g >= 1.0.

    Notes
    -----
    Implementation based on the `initseq` method from the `mcmc` R package by Geyer.

    References
    ----------
    [1] Geyer, CJ. Practical Markov chain Monte Carlo. Statistical Science 7(4):473-511, 1992.

    Examples
    --------

    Compute statistical inefficiency of timeseries data with known correlation time using different methods.

    >>> from pymbar import testsystems
    >>> A_n = testsystems.correlated_timeseries_example(N=100000, tau=10.0)
    >>> g_IPS = statisticalInefficiency_geyer(A_n, method='pos')
    >>> g_IMS = statisticalInefficiency_geyer(A_n, method='dec')
    >>> g_ICS = statisticalInefficiency_geyer(A_n, method='con')

    """

    if method not in ['pos', 'dec', 'con']:
        raise Exception("Unknown method '%s'; must be one of ['pos', 'dec', 'con']" % method)

    # Create numpy copies of input arguments.
    A_n = np.array(A_n)

    # Subtract off sample mean.
    A_n -= A_n.mean()

    # Get the length of the timeseries.
    N = A_n.size

    # Compute sample variance.
    gamma_zero = (A_n**2).sum() / N

    # Compute sequential covariance pairs.
    gamma_pos = list()
    for i in range(N/2):
        lag1 = 2*i
        gam1 = (A_n[0:(N-lag1)] * A_n[lag1:N]).sum() / N
        lag2 = lag1 + 1
        gam2 = (A_n[0:(N-lag2)] * A_n[lag2:N]).sum() / N

        # Terminate if sum is no longer positive.
        if (gam1 + gam2) < 0.0:
            break

        # Otherwise, store the consecutive sum.
        gamma_pos.append(gam1+gam2)

    # Number of nonnegative values in array.
    ngamma = len(gamma_pos)

    # Compute IPS gamma sequence.
    gamma_pos = np.array(gamma_pos)

    # Compute IMS gamma sequence.
    gamma_dec = np.array(gamma_pos)
    for i in range(ngamma - 1):
        if gamma_dec[i] < gamma_dec[i+1]:
            gamma_dec[i] = gamma_dec[i+1]

    # Compute ICS gamma sequence.
    gamma_con = np.array(gamma_dec)
    for i in range(ngamma-1, 0, -1):
        gamma_con[i] -= gamma_con[i-1]

    # Pool adjacent violators (PAVA) algorithm.
    puff = np.zeros([ngamma], np.float64)
    nuff = np.zeros([ngamma], np.int32)
    nstep = 0
    for j in range(1,ngamma):
        puff[nstep] = gamma_con[j]
        nuff[nstep] = 1
        nstep += 1
        while (nstep > 1) and ((puff[nstep-1] / nuff[nstep-1]) < (puff[nstep-2] / nuff[nstep-2])):
            puff[nstep-2] += puff[nstep-1]
            nuff[nstep-2] += nuff[nstep-1]
            nstep -= 1

    j = 1
    for jstep in range(nstep):
        muff = puff[jstep] / nuff[jstep]
        for k in range(nuff[jstep]):
            gamma_con[j] = gamma_con[j-1] + muff
            j += 1

    # Compute sample variance estimates.
    var_pos = (2 * gamma_pos.sum() - gamma_zero) / N
    var_dec = (2 * gamma_dec.sum() - gamma_zero) / N
    var_con = (2 * gamma_con.sum() - gamma_zero) / N

    # Compute statistical inefficiencies from sample mean var = var(A_n) / (N/g)
    # g = var / (var(A_n)/N)
    var_uncorr = gamma_zero / N
    g_pos = var_pos / var_uncorr
    g_dec = var_dec / var_uncorr
    g_con = var_con / var_uncorr

    # DEBUG
    #print "pos dec con : %12.3f %12.3f %12.3f" % (g_pos, g_dec, g_con)

    # Select appropriate g.
    if method == 'pos':
        g = g_pos
    elif method == 'dec':
        g = g_dec
    elif method == 'con':
        g = g_con

    # g must be at least unity
    if (g < 1.0):
        g = 1.0

    # Return the computed statistical inefficiency.
    return g

def detectEquilibration(A_t, nskip=1, method='multiscale'):
    """Automatically detect equilibrated region of a dataset using a heuristic that maximizes number of effectively uncorrelated samples.

    Parameters
    ----------
    A_t : np.ndarray
        timeseries
    nskip : int, optional, default=1
        number of samples to sparsify data by in order to speed equilibration detection
    method : str, optional, default='geyer'
        Method to use for computing statistical inefficiency; one of ['geyer', 'multiscale']

    Returns
    -------
    t : int
        start of equilibrated data
    g : float
        statistical inefficiency of equilibrated data
    Neff_max : float
        number of uncorrelated samples

    ToDo
    ----
    Consider implementing a binary search for Neff_max.

    Notes
    -----
    If your input consists of some period of equilibration followed by
    a constant sequence, this function treats the trailing constant sequence
    as having Neff = 1.

    Examples
    --------

    Determine start of equilibrated data for a correlated timeseries.

    >>> from pymbar import testsystems
    >>> A_t = testsystems.correlated_timeseries_example(N=1000, tau=10.0) # generate a test correlated timeseries
    >>> [t, g, Neff_max] = detectEquilibration(A_t) # compute indices of uncorrelated timeseries

    Determine start of equilibrated data for a correlated timeseries with a shift.

    >>> from pymbar import testsystems
    >>> A_t = testsystems.correlated_timeseries_example(N=1000, tau=10.0) + 2.0 # generate a test correlated timeseries
    >>> B_t = testsystems.correlated_timeseries_example(N=10000, tau=10.0) # generate a test correlated timeseries
    >>> C_t = np.concatenate([A_t, B_t])
    >>> [t, g, Neff_max] = detectEquilibration(C_t, nskip=50) # compute indices of uncorrelated timeseries

    """
    T = A_t.size

    # Special case if timeseries is constant.
    if A_t.std() == 0.0:
        return (0, 1, 1)  # Changed from Neff=N to Neff=1 after issue #122

    if method == 'multiscale':
        statisticalInefficiency = statisticalInefficiency_multiscale
    elif method == 'geyer':
        statisticalInefficiency = statisticalInefficiency_geyer
    else:
        raise Exception("Method must be one of ['multiscale', 'geyer']")

    g_t = np.ones([T - 1], np.float32)
    Neff_t = np.ones([T - 1], np.float32)
    for t in range(0, T - 1, nskip):
        g_t[t] = statisticalInefficiency(A_t[t:T])
        Neff_t[t] = (T - t + 1) / g_t[t]
    Neff_max = Neff_t.max()
    t = Neff_t.argmax()
    g = g_t[t]

    return (t, g, Neff_max)

max_stars_repo_name: jchodera/alchemical-best-practices

max_issues_repo_name: jchodera/alchemical-best-practices

max_forks_repo_name: jchodera/alchemical-best-practices


text:
import argparse
import collections
import logging
import math
import os
import re
import time
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp
import torch
import torch.nn as nn
import torch.utils.data
import transformers
from data_pipeline import create_pretraining_dataloader
from fp16 import FP16_Module, FP16_Optimizer, load_fp16_optimizer, save_fp16_optimizer
from learning_rates import AnnealingLR
from smdistributed.modelparallel.torch.nn import FusedLayerNorm as LayerNorm
from smdistributed.modelparallel.torch.nn.huggingface.gpt2 import (
    translate_hf_state_dict_to_smdistributed,
    translate_state_dict_to_hf_gpt2,
)
from torch import optim
from torch.nn.parallel.distributed import DistributedDataParallel
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_CAUSAL_LM_MAPPING,
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    GPT2Config,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import is_main_process

logger = logging.getLogger(__name__)

def get_learning_rate_scheduler(optimizer, args):

    # Add linear learning rate scheduler.
    if args.lr_decay_iters is not None:
        num_iters = args.lr_decay_iters
    else:
        num_iters = args.max_steps
    num_iters = max(1, num_iters)
    init_step = 0
    warmup_iter = args.warmup * num_iters
    plateau_iter = warmup_iter + args.plateau * num_iters
    lr_scheduler = AnnealingLR(
        optimizer,
        start_lr=args.lr,
        warmup_iter=warmup_iter,
        plateau_iter=plateau_iter,
        total_iters=num_iters,
        decay_style=args.lr_decay_style,
        last_iter=init_step,
        min_lr=args.min_lr,
        use_checkpoint_lr_scheduler=args.load_partial or args.load_full,
        override_lr_scheduler=False,
    )

    return lr_scheduler

def get_param_groups_by_weight_decay(module):
    weight_decay_params = {"params": []}
    no_weight_decay_params = {"params": [], "weight_decay": 0.0}
    param_ids = set()
    for module_ in module.modules():
        if isinstance(module_, LayerNorm):
            for p in list(module_._parameters.values()):
                if p is not None and id(p) not in param_ids:
                    no_weight_decay_params["params"].append(p)
                    param_ids.add(id(p))
        else:
            for n, p in list(module_._parameters.items()):
                if p is not None and n != "bias" and id(p) not in param_ids:
                    weight_decay_params["params"].append(p)
                    param_ids.add(id(p))
            for n, p in list(module_._parameters.items()):
                if p is not None and n == "bias" and id(p) not in param_ids:
                    no_weight_decay_params["params"].append(p)
                    param_ids.add(id(p))
    return weight_decay_params, no_weight_decay_params

# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    if args.logits_output:
        output = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
        loss = output["loss"]
    else:
        loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    if args.fp16:
        optimizer.backward(loss, update_master_grads=False)
    else:
        model.backward(loss)
    if args.logits_output:
        return output
    return loss

# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    return loss

def save_ckptsum(args, model, optimizer, filename):
    results = collections.defaultdict(dict)
    model_result = collections.defaultdict(dict)

    if args.fp16:
        from fp16.fp16util import register_optimizer_hooks

        register_optimizer_hooks(model)

    def _get_optimizer_result(optimizer_states):
        _optimizer_result = collections.defaultdict(dict)
        for param_idx, state in optimizer_states.items():
            for key, val in state.items():
                if isinstance(val, torch.Tensor):
                    _optimizer_result["tensors"][f"{param_idx}_{key}"] = torch.sum(val)
                else:
                    _optimizer_result["scalars"][f"{param_idx}_{key}"] = val
        return _optimizer_result

    if not args.shard_optimizer_state:
        optimizer_result = _get_optimizer_result(optimizer.local_state_dict()["state"])
    else:
        local_state_dict = optimizer.local_state_dict()["state"]
        if smp.rdp_rank() == 0:
            optimizer_result = []
            for partial_local_state_dict in local_state_dict:
                optimizer_result.append(_get_optimizer_result(partial_local_state_dict))

    for param_name, param in model.local_state_dict().items():
        if isinstance(param, torch.Tensor):
            model_result["tensors"][param_name] = torch.sum(param)
        else:
            model_result["scalars"][param_name] = param

    if smp.rdp_rank() == 0:
        results["optimizer"] = optimizer_result
        results["model"] = model_result
        smp.save(results, filename)

def load_and_verify_ckptsum(args, model, optimizer, filename):
    results = smp.load(filename)
    optimizer_result = (
        results["optimizer"]
        if not args.shard_optimizer_state
        else results["optimizer"][smp.rdp_rank()]
    )
    model_result = results["model"]

    def opt_check_fn(mod, opt):
        loaded_opt_states = (
            opt.orig_state_dict()["state"]
            if args.shard_optimizer_state
            else opt.local_state_dict()["state"]
        )
        for param_idx, state in loaded_opt_states.items():
            for key, val in state.items():
                if isinstance(val, torch.Tensor):
                    assert torch.isclose(
                        torch.sum(val), optimizer_result["tensors"][f"{param_idx}_{key}"]
                    ), f"mismatch for param_idx: {param_idx}, key is {key}"
                else:
                    assert (
                        val == optimizer_result["scalars"][f"{param_idx}_{key}"]
                    ), f"mismatch for param_idx: {param_idx}, key is {key}"
        print("Optimizer save/load check passed successfully")

    def model_check_fn(mod, opt):
        for param_name, param in mod.local_state_dict().items():
            if isinstance(param, torch.Tensor):
                assert torch.isclose(
                    torch.sum(param), model_result["tensors"][param_name]
                ), f"mismatch for param_name: {param_name}"
            else:
                assert (
                    param == model_result["scalars"][param_name]
                ), f"mismatch for param_name: {param_name}"
        print("Model save/load check passed successfully")

    model.register_post_partition_hook(model_check_fn)
    model.register_post_step_hook(opt_check_fn)

def save(
    output_save_file,
    model,
    optimizer,
    lr_scheduler,
    model_config,
    num_params,
    total_steps,
    curr_train_path_index,
    args,
    partial=True,
    translate_to_hf=False,
    seq_length=1024,
    batch_idx=0,
):
    save_fn = save_fp16_optimizer
    save_dict = {
        "cli_args": args.__dict__,
        "num_params": num_params,
        "total_steps": total_steps,
        "curr_train_path_index": curr_train_path_index,
        "model_config": model_config,
        "batch_idx": batch_idx,
    }

    if lr_scheduler is not None:
        save_dict["lr_scheduler"] = lr_scheduler.state_dict()
    if partial:
        save_dict["model"] = model.local_state_dict()
    else:
        model_state_dict = model.state_dict(gather_to_rank0=True)
        if smp.rank() == 0:
            save_dict["model"] = (
                translate_state_dict_to_hf_gpt2(model_state_dict, seq_length)
                if translate_to_hf
                else model_state_dict
            )

    if args.fp16:
        if not partial and args.skip_full_optimizer:
            print("Skipping saving the final optimizer state")
        else:
            if args.shard_optimizer_state == 0 or partial:
                save_dict["optimizer"] = save_fn(args, model, optimizer, partial=partial)
            else:
                print(
                    "Saving the full optimizer state does not work with shard_optimizer_state > 0! Skipping..."
                )
    else:
        # fp32
        if partial:
            save_dict["optimizer"] = optimizer.local_state_dict()
        else:
            if not args.skip_full_optimizer:
                save_dict["optimizer"] = optimizer.state_dict()
            else:
                print("Skipping saving of full optimizer state")

    if (smp.rdp_rank() == 0 and partial) or smp.rank() == 0:
        smp.save(save_dict, output_save_file, partial=partial)

    print(f"Finished checkpointing after {total_steps} steps: {output_save_file}")

def load_model_and_optimizer(
    output_dir,
    model,
    optimizer,
    lr_scheduler,
    partial,
    args,
    translate_from_hf=False,
    seq_length=1024,
    load_model=True,
    load_optimizer=True,
    num_params=0,
):
    # Find longest-trained checkpoint
    re_pattern = f"trained_gpt_nparams-{num_params}_steps-(?P<total_steps>\d+)\.pt"
    if partial:
        re_pattern += "_(?P<rank>\d+)"
    else:
        re_pattern += "$"

    ckpt_paths = sorted(
        [
            (int(re.match(re_pattern, p).group("total_steps")), os.path.join(output_dir, p))
            for p in os.listdir(output_dir)
            if re.match(re_pattern, p)
        ],
        reverse=True,
    )
    if not ckpt_paths:
        raise Exception(
            f'No checkpoints could be found in "{output_dir}".  Candidates: {os.listdir(output_dir)}'
        )

    local_ckpt_path = ckpt_paths[0][1]

    if partial:
        # need to pass prefix without ranks to smp
        local_ckpt_path = local_ckpt_path.split(".pt")[0] + ".pt"

    checkpoint = smp.load(local_ckpt_path, partial=partial)

    if load_model:
        checkpointed_model = (
            translate_hf_state_dict_to_smdistributed(checkpoint["model"], seq_length)
            if translate_from_hf
            else checkpoint["model"]
        )
        model.load_state_dict(checkpointed_model, same_partition_load=args.same_partition_load > 0)
        if lr_scheduler is not None:
            lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])

    if load_optimizer:
        # Loading loss scale eagerly
        opt_state_dict = checkpoint["optimizer"]
        optimizer.loss_scaler = opt_state_dict["loss_scaler"]
        optimizer.loss_scaler.model = model
        optimizer.dynamic_loss_scale = opt_state_dict["dynamic_loss_scale"]
        optimizer.overflow = opt_state_dict["overflow"]
        optimizer.first_closure_call_this_step = opt_state_dict["first_closure_call_this_step"]

        def opt_load_hook(mod, opt):
            load_fn = load_fp16_optimizer
            if args.fp16:
                if not partial and args.skip_full_optimizer:
                    print(
                        "Skipping loading the final optimizer state, and reloading master_params from model_params"
                    )
                    opt.reload_model_params()
                else:
                    load_fn(args, mod, opt, checkpoint, partial=partial)
            else:
                # fp32
                if not partial and args.skip_full_optimizer:
                    print("Skipping loading the final optimizer state")
                else:
                    opt.load_state_dict(checkpoint["optimizer"])

        model.register_post_step_hook(opt_load_hook)

    print(f'Loaded model from "{local_ckpt_path}"')

    batch_idx = 0
    if "batch_idx" in checkpoint:
        batch_idx = checkpoint["batch_idx"]

    return (
        model,
        optimizer,
        checkpoint["total_steps"],
        checkpoint["curr_train_path_index"],
        batch_idx,
    )

def delete_oldest_ckpt(args):
    re_pattern = "trained_gpt_nparams-(?P<num_params>\d+)_steps-(?P<total_steps>\d+)\.pt"

    # partial
    re_pattern += "_(?P<pp_rank>\d+)_(?P<tp_rank>\d+)"

    paths_per_step = collections.defaultdict(list)

    for p in os.listdir(args.checkpoint_dir):
        if re.match(re_pattern, p):
            step = int(re.match(re_pattern, p).group("total_steps"))
            path = os.path.join(args.checkpoint_dir, p)
            paths_per_step[step].append(path)

    if paths_per_step:
        oldest_step = sorted(paths_per_step.keys())[0]
        num_parts = len(paths_per_step[oldest_step])
        if len(paths_per_step) > args.num_kept_checkpoints:
            # delete oldest step
            for p in paths_per_step[oldest_step]:
                os.remove(p)
    # else We still haven't reached maximum number of checkpoints -- no need to delete older ones
    return None

def eval_model(model, dataloader, num_batches, use_wiki_data):
    model = model.eval()
    n_batches = 0
    loss = 0.0

    with torch.no_grad():
        for batch_idx, input_data in enumerate(dataloader):
            if use_wiki_data:
                input_ids, _, attention_mask, _, _ = input_data
            else:
                input_ids, attention_mask = input_data
            if batch_idx >= num_batches:
                break

            loss += test_step(model, input_ids, attention_mask).reduce_mean()
            n_batches += 1

    if n_batches > 0:
        torch.distributed.all_reduce(loss, group=smp.get_dp_process_group())
        loss /= smp.dp_size()
        loss /= n_batches
        loss = loss.item()
        ppl = math.exp(loss)
    else:
        loss = -1.0
        ppl = -1.0

    return loss, ppl

def train(
    model,
    optimizer,
    lr_scheduler,
    model_config,
    start_train_path_index,
    start_batch_index,
    num_params,
    total_steps,
    args,
):
    model.train()
    if args.parallel_proc_data_processing:
        pool = ProcessPoolExecutor(1)

    dp_rank = smp.dp_rank() if not args.prescaled_batch else smp.rdp_rank()
    dp_size = smp.dp_size() if not args.prescaled_batch else smp.rdp_size()
    data_type = "wiki" if args.use_wiki_data else "openwebtext"

    if args.use_wiki_data:
        train_paths = sorted(
            [
                os.path.join(args.training_dir, p)
                for p in os.listdir(args.training_dir)
                if os.path.isfile(os.path.join(args.training_dir, p)) and "training" in p
            ]
        )
    else:
        if args.zipped_data > 0:
            file_extension = ".json.gz"
        else:
            file_extension = ".json"
        train_paths = sorted(
            [
                os.path.join(args.training_dir, p)
                for p in os.listdir(args.training_dir)
                if p.endswith(file_extension)
            ]
        )

    train_dataloader = create_pretraining_dataloader(
        [train_paths[start_train_path_index]],
        args.train_batch_size,
        args.max_context_width,
        seed=args.seed,
        dp_rank=dp_rank,
        dp_size=dp_size,
        shuffle=args.same_seed < 1,
        zipped=args.zipped_data > 0,
        use_last_file_only=args.fast_validation > 0,
        data_type=data_type,
    )

    if args.validation_freq is not None:
        # load all validation examples
        if smp.rank() == 0:
            print("Creating val dataloader")
        if args.use_wiki_data:
            val_paths = sorted(
                [
                    os.path.join(args.test_dir, p)
                    for p in os.listdir(args.test_dir)
                    if os.path.isfile(os.path.join(args.test_dir, p)) and "testing" in p
                ]
            )

        else:
            if args.zipped_data > 0:
                file_extension = ".json.gz"
            else:
                file_extension = ".json"
            val_paths = sorted(
                [
                    os.path.join(args.test_dir, p)
                    for p in os.listdir(args.test_dir)
                    if p.endswith(file_extension)
                ]
            )
        val_dataloader = create_pretraining_dataloader(
            val_paths,
            args.val_batch_size,
            args.max_context_width,
            seed=args.seed,
            dp_rank=dp_rank,
            dp_size=dp_size,
            shuffle=True,
            zipped=args.zipped_data > 0,
            use_last_file_only=args.fast_validation > 0,
            data_type=data_type,
        )
        if smp.rank() == 0:
            print("Created val dataloader")

    start = time.time()
    throughput = None
    to_save = {"loss": [], "val_loss": []}
    loss_metric = 0

    def should_record():
        # only record the ranks that in the tp group that contains global rank 0
        if smp.tp_size() > 1:
            tp_group = smp.get_tp_group()
            return 0 in tp_group
        else:
            return smp.rank() == 0

    # Set the same seed for computation
    set_seed(args.seed)

    for index in range(start_train_path_index, args.epochs * len(train_paths)):
        next_train_path_index = (index + 1) % len(train_paths)
        curr_train_path_index = index % len(train_paths)

        if total_steps >= args.max_steps:
            break

        if args.parallel_proc_data_processing:
            dataset_future = pool.submit(
                create_pretraining_dataloader,
                [train_paths[next_train_path_index]],
                args.train_batch_size,
                args.max_context_width,
                seed=args.seed,
                dp_rank=dp_rank,
                dp_size=dp_size,
                shuffle=args.same_seed < 1,
                zipped=args.zipped_data > 0,
                use_last_file_only=args.fast_validation > 0,
                data_type=data_type,
            )

        if smp.rank() == 0:
            if args.use_wiki_data:
                print(f"Reading data from training path {train_dataloader.dataset.input_file}")
            else:
                print(f"Reading data from training path {train_dataloader.dataset.input_paths}")

        for batch_idx, input_data in enumerate(train_dataloader):
            if batch_idx < start_batch_index:
                if smp.rank() == 0:
                    print(
                        f"Resuming from saved batch index {start_batch_index}, skipping batch {batch_idx}..."
                    )
                continue
            else:
                start_batch_index = 0

            if args.use_wiki_data:
                input_ids, _, attention_mask, _, _ = input_data
            else:
                input_ids, attention_mask = input_data

            if total_steps >= args.max_steps:
                break

            step_start = time.time()

            if args.fp16:
                optimizer.zero_grad(set_grads_to_None=True)
            else:
                optimizer.zero_grad()

            if args.logits_output:
                train_output = train_step(model, optimizer, input_ids, attention_mask, args)
                loss_mb = train_output["loss"]
                logits_mb = train_output["logits"]
                if smp.tp_size() > 1:
                    logits = torch.cat(tuple(logits_mb.outputs), dim=1)
                else:
                    logits = torch.cat(tuple(logits_mb.outputs), dim=0)
            else:
                # Return value, loss_mb is a StepOutput object
                loss_mb = train_step(model, optimizer, input_ids, attention_mask, args)

            # smdistributed: Average the loss across microbatches.
            loss = loss_mb.reduce_mean()
            if not args.validation_freq:
                loss_metric = loss.item()

            if args.fp16:
                optimizer.update_master_grads()
                optimizer.clip_master_grads(args.grad_clip)
                optimizer.step()
                overflow = optimizer.overflow
            else:
                optimizer.step()
            if not (args.fp16 and overflow):
                lr_scheduler.step()

            total_steps += 1
            time_elapsed = time.time() - start
            step_time = time.time() - step_start
            sample_processed = input_ids.shape[0] * dp_size
            throughput = sample_processed / step_time
            if smp.rank() == 0 and not total_steps % args.logging_freq:
                print(
                    f"({int(time_elapsed)}s), Batch {total_steps - 1} Loss: {loss.item()}, Speed: {throughput} samples/sec"
                )

            # evaluate on validation
            if args.validation_freq and not (total_steps % args.validation_freq):
                cur_state = np.random.get_state()
                model = model.eval()
                val_loss, val_ppl = eval_model(
                    model, val_dataloader, args.validation_batches, args.use_wiki_data
                )
                if is_main_process(smp.rank()):
                    print(
                        f"({int(time.time()-start)}s) Batch {total_steps - 1} Validation loss: {val_loss}"
                    )
                    print(
                        f"({int(time.time()-start)}s) Batch {total_steps - 1} Validation perplexity: {val_ppl}"
                    )
                loss_metric = val_loss
                if args.logits_output:
                    to_save["val_loss"].append(val_loss)
                model = model.train()
                if args.preserve_np_state > 0:
                    np.random.set_state(cur_state)

            # checkpoint
            if not (total_steps % args.checkpoint_freq):
                base_path = f"trained_gpt_nparams-{num_params}_steps-{total_steps}.pt"
                out_path = os.path.join(args.checkpoint_dir, base_path)
                total_ckpts = total_steps // args.checkpoint_freq

                # save_or_verify_ckptsum if this is the last checkpoint
                if (args.save_or_verify_ckptsum and total_steps >= args.max_steps) or (
                    (total_ckpts + 1) * args.checkpoint_freq
                ) > args.max_steps:
                    # Save optimizer and model tensor sums and scalars before saving
                    save_ckptsum(
                        args,
                        model,
                        optimizer,
                        filename=os.path.join(args.model_dir, "saved_partial_sum"),
                    )

                save(
                    out_path,
                    model,
                    optimizer,
                    lr_scheduler,
                    model_config,
                    num_params,
                    total_steps,
                    curr_train_path_index,
                    args,
                    partial=True,
                    batch_idx=batch_idx + 1,
                )

                if smp.local_rank() == 0:
                    delete_oldest_ckpt(args)

            if args.logits_output:
                to_save["loss"].append(loss.item())

        if total_steps >= args.max_steps:
            if should_record() and args.logits_output:
                to_save["logits"] = logits.detach().cpu()
                output_file = f"rank_{smp.rank()}_" + args.logits_output
                torch.save(to_save, os.path.join(args.model_dir, output_file))
                print(f"logits and loss saved at {os.path.join(args.model_dir, output_file)}")
            break

        del train_dataloader

        if args.parallel_proc_data_processing:
            s = time.time()
            train_dataloader = dataset_future.result(timeout=None)
            wait_time = time.time() - s
            if wait_time > 1:
                print(
                    f"[{smp.rank()}] Waited {wait_time} for data loader to be ready. Please check if dataloader performance can be improved to avoid these waits."
                )
        else:
            train_dataloader = create_pretraining_dataloader(
                [train_paths[next_train_path_index]],
                args.train_batch_size,
                args.max_context_width,
                seed=args.seed,
                dp_rank=dp_rank,
                dp_size=dp_size,
                shuffle=args.same_seed < 1,
                zipped=args.zipped_data > 0,
                use_last_file_only=args.fast_validation > 0,
                data_type=data_type,
            )

    return total_steps, throughput, loss_metric

def parse_args():
    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.

    opt_grp = parser.add_argument_group(
        title="optimization", description="arguments for optimization"
    )
    opt_grp.add_argument(
        "--train_batch_size",
        type=int,
        default=4,
        help="batch size per dp rank, for tensor parallelism degree 8 with pipeline parallel degree 1 this means 8*this batch size per node",
    )
    opt_grp.add_argument("--val_batch_size", type=int, default=4)
    opt_grp.add_argument("--max_steps", type=int, default=5000)
    opt_grp.add_argument("--seed", type=int, default=12345)
    opt_grp.add_argument("--same_seed", type=int, default=0)
    opt_grp.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    opt_grp.add_argument("--fp16", default=0, type=int, help="automatic mixed precision training")
    opt_grp.add_argument(
        "--fp32_grad_accumulation", default=0, type=int, help="Enable FP32 Grad accumulation"
    )
    opt_grp.add_argument("--grad_clip", default=1.0, type=float, help="gradient clipping")
    opt_grp.add_argument("--weight_decay", default=0.01, type=float, help="weight decay")
    opt_grp.add_argument(
        "--beta1", default=0.9, type=float, help="beta1 parameter for Adam optimizer"
    )
    opt_grp.add_argument(
        "--beta2", default=0.95, type=float, help="beta2 parameter for Adam optimizer"
    )
    opt_grp.add_argument(
        "--activation_checkpointing",
        type=int,
        default=1,
        help="enable gradient checkpointing to reduce memory consumption",
    )
    parser.add_argument(
        "--logging_freq", type=int, default=1, help="number of iterations between logging"
    )

    # I/O
    io_grp = parser.add_argument_group(title="io", description="location for input and output")
    io_grp.add_argument("--use_wiki_data", type=int, default=0, help="use wiki corpus data for training")
    io_grp.add_argument("--zipped_data", type=int, default=1, help="input data is zipped files")
    io_grp.add_argument(
        "--epochs", type=int, default=3, help="times of iterating over the training dataset"
    )
    io_grp.add_argument("--output-data-dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    io_grp.add_argument(
        "--checkpoint-dir",
        type=str,
        default="/opt/ml/checkpoints",
        help="Saves partial checkpoints (model, optimizer) to this dir, and loads latest checkpoint from this if load_partial is specified.",
    )
    io_grp.add_argument(
        "--model-dir",
        type=str,
        default=os.environ["SM_MODEL_DIR"],
        help="Saves full model for inference to this dir. Also used if load_full is given to load the model. Note the lack of optimizer state here.",
    )
    io_grp.add_argument("--training-dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    io_grp.add_argument("--test-dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
    io_grp.add_argument(
        "--parallel_proc_data_processing",
        type=int,
        default=0,
        help="Load data in parallel with a different process. At any point a process can have two files in memory. With tensor parallelism, each of the 8 processes on an instance will then have 2 files in memory. Depending on file sizes this may or may not be feasible. With pipeline parallelism this was not a problem as only 1 rank on an instance loaded data.",
    )
    io_grp.add_argument(
        "--save_final_full_model",
        type=int,
        default=0,
        help="Enabling this will save a combined model only at the end",
    )
    io_grp.add_argument(
        "--skip_full_optimizer",
        type=int,
        default=1,
        help="Disabling this will also save the full optimizer state",
    )
    io_grp.add_argument("--load_partial", type=int, default=0, help="Load from partial checkpoints")
    io_grp.add_argument("--load_full", type=int, default=0, help="Load from full checkpoints")
    io_grp.add_argument(
        "--logits_output", type=str, default="", help="Path to save logits and loss"
    )
    io_grp.add_argument("--prescaled_batch", type=int, default=1, help="use prescaled batch")

    # configure model size
    model_grp = parser.add_argument_group(
        title="model", description="arguments to describe model configuration"
    )
    model_grp.add_argument("--max_context_width", type=int, default=1024)
    model_grp.add_argument("--hidden_width", type=int, default=768)
    model_grp.add_argument("--num_layers", type=int, default=12)
    model_grp.add_argument("--num_heads", type=int, default=12)
    model_grp.add_argument("--resid_pdrop", type=float, default=0.1)
    model_grp.add_argument("--embd_pdrop", type=float, default=0.1)
    model_grp.add_argument("--attn_pdrop", type=float, default=0.1)
    model_grp.add_argument("--summary_first_pdrop", type=float, default=0.1)
    model_grp.add_argument("--use_adamw", type=int, default=0, help="Use adamw optimizer")

    smp_grp = parser.add_argument_group(title="smp", description="smp")
    smp_grp.add_argument("--tensor_parallel_degree", type=int, default=8)
    smp_grp.add_argument("--pipeline_parallel_degree", type=int, default=1)
    smp_grp.add_argument("--microbatches", type=int, default=1)
    smp_grp.add_argument("--active_microbatches", type=int, default=None)
    smp_grp.add_argument("--optimize", type=str, default="speed")
    smp_grp.add_argument("--activation_strategy", type=str, default="each")
    smp_grp.add_argument("--shard_optimizer_state", type=int, default=0)
    smp_grp.add_argument("--offload_activations", type=int, default=0)
    smp_grp.add_argument("--fast_mode", type=int, default=0)
    smp_grp.add_argument("--static_mode", type=int, default=0)
    smp_grp.add_argument("--delayed_param", type=int, default=0)
    smp_grp.add_argument("--same_partition_load", type=int, default=0)
    smp_grp.add_argument("--attention_in_fp32", type=int, default=1)

    parser.add_argument(
        "--num_kept_checkpoints",
        type=int,
        default=5,
        help="how many checkpoints to keep before deleting",
    )
    parser.add_argument(
        "--checkpoint_freq",
        type=int,
        default=10000,
        help="number of iterations between checkpointing",
    )
    parser.add_argument(
        "--validation_freq",
        type=int,
        default=None,
        help="number of iterations to print validation loss",
    )
    parser.add_argument(
        "--validation_batches",
        type=int,
        default=10,
        help="number of batches to estimate validation loss",
    )
    parser.add_argument(
        "--manual_partition",
        type=int,
        default=0,
        help="evenly distribute layers across the partitions",
    )
    parser.add_argument(
        "--match_weights", type=int, default=0, help="Get weights from the original model"
    )
    parser.add_argument(
        "--preserve_np_state",
        type=int,
        default=0,
        help="Perserve the numpy random state between validation",
    )
    parser.add_argument(
        "--fast_validation",
        type=int,
        default=1,
        help="Running validation only with the last data file for faster speed",
    )

    # learning rate
    lr_grp = parser.add_argument_group(
        title="lr", description="arguments for learning rate schedule"
    )
    lr_grp.add_argument("--lr", type=float, default=None, help="Initial learning rate.")
    lr_grp.add_argument(
        "--lr_decay_style",
        type=str,
        default="linear",
        choices=["constant", "linear", "cosine", "exponential", "plateau"],
        help="Learning rate decay function.",
    )
    lr_grp.add_argument(
        "--lr_decay_iters",
        type=int,
        default=None,
        help="number of iterations to decay learning rate over," " If None defaults to train iters",
    )
    lr_grp.add_argument(
        "--min_lr",
        type=float,
        default=0.0,
        help="Minumum value for learning rate. The scheduler" "clip values below this threshold.",
    )
    lr_grp.add_argument(
        "--warmup",
        type=float,
        default=0.01,
        help="Percentage of total iterations to warmup on "
        "(.01 = 1 percent of all training iters).",
    )
    lr_grp.add_argument(
        "--plateau",
        type=float,
        default=0.4,
        help="Percentage of total iterations to keep at max if using plateau lr",
    )

    ci_grp = parser.add_argument_group(title="ci", description="ci related settings")
    ci_grp.add_argument("--ci", default=False, action="store_true", help="Whether enable ci")
    ci_grp.add_argument("--time_to_train", type=int, help="time to train threshold")
    ci_grp.add_argument("--throughput", type=float, help="throughput threshold")
    ci_grp.add_argument("--loss", type=float, help="loss threshold")
    ci_grp.add_argument(
        "--save_or_verify_ckptsum", default=False, action="store_true", help="Whether to save sum"
    )

    args, _ = parser.parse_known_args()
    return args

def main():
    args = parse_args()

    if args.shard_optimizer_state > 0 and not args.skip_full_optimizer:
        raise ValueError(
            "If shard_optimizer_state is enabled, skip_full_optimizer must also be enabled. Full optimizer saving is currently not supported under optimizer state sharding."
        )

    # any value here is overriden by the config set in notebook when launching the sagemaker job
    smp_config = {
        "ddp": True,
        "tensor_parallel_degree": args.tensor_parallel_degree,
        "pipeline_parallel_degree": args.pipeline_parallel_degree,
        "microbatches": args.microbatches,
        # if activation_checkpointing true checkpoints transformer layers below
        "checkpoint_attentions": False if args.activation_checkpointing else True,
        "shard_optimizer_state": args.shard_optimizer_state > 0,
        "prescaled_batch": args.prescaled_batch > 0,
        "_match_weights": args.match_weights > 0,
        "fp16_params": args.fp16 > 0,
        "offload_activations": args.offload_activations > 0,
        "optimize": args.optimize,
        "auto_partition": False if args.manual_partition else True,
        "default_partition": 0,
        "_fp32_grad_accumulation": args.fp32_grad_accumulation > 0,
        "static_mode": args.static_mode > 0,
        "fast_mode": args.fast_mode > 0,
    }
    if args.active_microbatches is not None:
        smp_config["active_microbatches"] = args.active_microbatches

    smp.init(smp_config)

    if smp.rank() == 0:
        print("Arguments:", args.__dict__)
        print(f"Transformers version: {transformers.__version__}")
        print(f"smdistributed.modelparallel version: {smdistributed.modelparallel.__version__}")
        print(f"smdistributed config: {smp_config}")

    if args.save_final_full_model and smp.rank() == 0:
        print(
            f"[Warning] Note that save_final_full_model only saves the final model at the end of all steps. It does not save optimizer state. Optimizer state is only saved with partial models which are saved at checkpointing_freq during training. If you want to restart training you need partial checkpoints."
        )

    if smp.local_rank() == 0:
        for path in [args.model_dir, args.checkpoint_dir]:
            if not os.path.exists(path):
                os.makedirs(path, exist_ok=True)

    model_config = GPT2Config(
        vocab_size=50257,
        n_positions=args.max_context_width,
        n_ctx=args.max_context_width,
        n_embd=args.hidden_width,
        n_layer=args.num_layers,
        n_head=args.num_heads,
        n_inner=None,
        activation_function="gelu_new",
        resid_pdrop=args.resid_pdrop,
        embd_pdrop=args.embd_pdrop,
        attn_pdrop=args.attn_pdrop,
        layer_norm_epsilon=1e-05,
        initializer_range=0.02,
        summary_type="cls_index",
        summary_use_proj=True,
        summary_activation=None,
        summary_proj_to_labels=True,
        summary_first_dropout=args.summary_first_pdrop,
        # gradient_checkpointing=args.gradient_checkpointing > 0,
        use_cache=True,
        bos_token_id=50256,
        eos_token_id=50256,
        return_dict=True,
    )

    # the following improves start-up time by skipping proper initialization
    # of weights in the original model. this is not a problem because DistributedModel
    # will override those weights anyway when tensor_parallel_degree > 1.
    if smp.tp_size() > 1 and args.match_weights < 1:
        from transformers.modeling_utils import PreTrainedModel

        PreTrainedModel.init_weights = lambda x: None

    set_seed(args.seed)

    if args.fp16:
        torch.set_default_dtype(torch.float16)
    with smp.tensor_parallelism(
        enabled=smp.tp_size() > 1, attention_in_fp32=args.attention_in_fp32 > 0
    ):
        with smp.delay_param_initialization(
            enabled=(smp.tp_size() > 1 and args.match_weights < 1 and args.delayed_param > 0)
        ):
            model = AutoModelForCausalLM.from_config(model_config)
    torch.set_default_dtype(torch.float32)

    if args.fp16:
        model = FP16_Module(model)

    num_params = sum([np.prod(p.size()) for p in model.parameters()])
    if smp.rank() == 0:
        print(f"# total parameters: {num_params}")

    # smdistributed: Set the device to the GPU ID used by the current process.
    # Input tensors should be transferred to this device.
    torch.cuda.set_device(smp.local_rank())
    device = torch.device("cuda")

    if not args.same_seed:
        # Set seed by tp_rank to prevent weights from being the same on different tp_ranks
        set_seed(args.seed + smp.tp_rank())

    # smdistributed: Use the DistributedModel container to provide the model
    # to be partitioned across different ranks. For the rest of the script,
    # the returned DistributedModel object should be used in place of
    # the model provided for DistributedModel class instantiation.
    if args.fp16:
        torch.set_default_dtype(torch.float16)
    model = smp.DistributedModel(model, trace_device="gpu")

    if args.fp16:
        m = model.module
    else:
        m = model

    if smp.tp_size() > 1:
        transformer_layers = m.module.module.transformer.seq_layers
    else:
        transformer_layers = m.module.module.transformer.h

    if args.manual_partition:
        print(f"Manual partition enabled")
        # evenly distribute layers across all partitions
        div, rem = divmod(args.num_layers, smp.pp_size())
        get_num_layers = lambda x: (div + 1 if x >= smp.pp_size() - rem else div)
        assignments = []
        for pp_rank in range(smp.pp_size()):
            nl = get_num_layers(pp_rank)
            print(f"{nl} layers assigned to partition {pp_rank}")
            assignments += [pp_rank for _ in range(nl)]

        for i, c in enumerate(transformer_layers.children()):
            smp.set_partition(c, assignments[i])

    torch.set_default_dtype(torch.float32)

    iter_model = model
    # Build parameter groups (weight decay and non-decay).
    while isinstance(iter_model, (DistributedDataParallel, FP16_Module)):
        iter_model = iter_model.module
    param_groups = get_param_groups_by_weight_decay(iter_model)

    if args.use_adamw > 0:
        optimizer = optim.AdamW(
            param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
        )
    else:
        optimizer = optim.Adam(
            param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
        )

    if args.activation_checkpointing:
        kwargs = {}
        if isinstance(transformer_layers, nn.Sequential):
            kwargs["pack_args_as_tuple"] = True
            kwargs["strategy"] = args.activation_strategy
        smp.set_activation_checkpointing(transformer_layers, **kwargs)

    if args.fp16:
        optimizer = FP16_Optimizer(
            model,
            optimizer,
            static_loss_scale=None,
            dynamic_loss_scale=True,
            use_smp=True,
            dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
            params_have_main_grad=args.fp32_grad_accumulation > 0,
            shard_optimizer_state=args.shard_optimizer_state > 0,
        )

    optimizer = smp.DistributedOptimizer(optimizer)
    lr_scheduler = get_learning_rate_scheduler(optimizer, args)

    if args.fp16:
        model.register_post_step_hook(lambda model, optimizer: optimizer.init_master_params())

    # load after wrapping model and optimizer with smp Distributed...
    if args.load_full or args.load_partial:
        if args.load_partial and args.load_full:
            print(
                "Since both --load_partial and --load_full set, will try to load from full checkpoint."
                "If the intention is to load from partial checkpoint, please don't set --load_full"
            )
        partial = not args.load_full
        path = args.checkpoint_dir if partial else args.model_dir
        translate_from_hf = not partial
        model, optimizer, total_steps, start_train_path_index, start_batch_index = load_model_and_optimizer(
            path,
            model,
            optimizer,
            lr_scheduler,
            partial,
            args,
            translate_from_hf=translate_from_hf,
            seq_length=args.max_context_width,
            load_model=True,
            load_optimizer=args.load_partial > 0,
            num_params=num_params,
        )
        if args.save_or_verify_ckptsum:
            filename = "saved_sum" if args.load_full else "saved_partial_sum"
            load_and_verify_ckptsum(
                args, model, optimizer, filename=os.path.join(args.model_dir, filename)
            )
    else:
        total_steps = 0
        start_train_path_index = 0
        start_batch_index = 0

    start = time.time()
    total_steps, throughput, loss = train(
        model,
        optimizer,
        lr_scheduler,
        model_config,
        start_train_path_index,
        start_batch_index,
        num_params,
        total_steps,
        args,
    )
    time_to_train = time.time() - start
    if args.ci:
        print(f"[SMP_METRIC]__GPT2__Time_to_train__{time_to_train}")
        print(f"[SMP_METRIC]__GPT2__samples/second__{throughput}")
        print(f"[SMP_METRIC]__GPT2__Loss__{loss}")
        if not args.load_partial and not args.load_full:
            assert time_to_train < args.time_to_train
            assert throughput > args.throughput
            if args.loss:
                assert loss < args.loss

    if args.save_final_full_model:
        # saves full model at the end

        base_path = f"trained_gpt_nparams-{num_params}_steps-{total_steps}.pt"
        out_path = os.path.join(args.model_dir, base_path)
        if args.save_or_verify_ckptsum:
            # Save optimizer and model tensor sums and scalars before saving
            save_ckptsum(args, model, optimizer, filename=os.path.join(args.model_dir, "saved_sum"))

        if smp.rdp_rank() == 0:
            save(
                out_path,
                model,
                optimizer,
                lr_scheduler,
                model_config,
                num_params,
                total_steps,
                -1,
                args,
                partial=False,
                translate_to_hf=smp.tp_size() > 1,
                seq_length=args.max_context_width,
            )

    smp.barrier()
    if smp.rank() == 0:
        print("SMP training finished successfully")

if __name__ == "__main__":
    main()

max_stars_repo_name: pollyrolly/amazon-sagemaker-examples

max_issues_repo_name: pollyrolly/amazon-sagemaker-examples

max_forks_repo_name: pollyrolly/amazon-sagemaker-examples