<commit_before>from __future__ import absolute_import

import numpy as np
import sys
from itertools import tee
from keras.models import Sequential, model_from_yaml
from pyspark.rdd import RDD

from .utils.functional_utils import add_params, get_neutral, divide_by

class SparkModel(object):
    def __init__(self, sc, master_network, nb_epoch=10, batch_size=32, verbose=0, validation_split=0.1):
        self.spark_context = sc
        self.master_network = master_network
        self.nb_epoch = nb_epoch
        self.batch_size = batch_size
        self.verbose = verbose
        self.validation_split = 0.1

    def get_train_config(self):
        train_config = {}
        train_config['nb_epoch'] = self.nb_epoch
        train_config['batch_size'] = self.batch_size
        train_config['verbose'] = self.verbose
        train_config['validation_split'] = self.validation_split
        return train_config

    def get_network(self):
        return self.master_network

    def set_network(self, network):
        self.master_network = network

    def train(self, rdd):
        num_partitions = rdd.getNumPartitions()
        yaml = self.master_network.to_yaml()
        parameters = self.spark_context.broadcast(self.master_network.get_weights())
        train_config = self.get_train_config()

        worker = SparkWorker(yaml, parameters, train_config)
        results = rdd.mapPartitions(worker.train)

        null_element = get_neutral(results.first())
        new_parameters = divide_by(results.fold(null_element, add_params), num_partitions)

        self.master_network.set_weights(new_parameters)

class SparkWorker(object):<commit_msg>Pass train parameters to train method<commit_after>from __future__ import absolute_import

import numpy as np
import sys
from itertools import tee
from keras.models import Sequential, model_from_yaml
from pyspark.rdd import RDD

from .utils.functional_utils import add_params, get_neutral, divide_by

def get_train_config(nb_epoch, batch_size, verbose, validation_split):
    train_config = {}
    train_config['nb_epoch'] = nb_epoch
    train_config['batch_size'] = batch_size
    train_config['verbose'] = verbose
    train_config['validation_split'] = validation_split
    return train_config


class SparkModel(object):
    def __init__(self, sc, master_network):
        self.spark_context = sc
        self.master_network = master_network

    def get_network(self):
        return self.master_network

    def set_network(self, network):
        self.master_network = network

    def train(self, rdd, nb_epoch=10, batch_size=32, verbose=0, validation_split=0.1):
        num_partitions = rdd.getNumPartitions()
        yaml = self.master_network.to_yaml()
        parameters = self.spark_context.broadcast(self.master_network.get_weights())
        train_config = get_train_config(nb_epoch, batch_size, verbose, validation_split)

        worker = SparkWorker(yaml, parameters, train_config)
        results = rdd.mapPartitions(worker.train)

        null_element = get_neutral(results.first())
        new_parameters = divide_by(results.fold(null_element, add_params), num_partitions)

        self.master_network.set_weights(new_parameters)

class SparkWorker(object):
{'fraction_non_alphanumeric': 0.08232745280099041, 'fraction_numerical': 0.004952027236149799, 'mean_word_length': 3.874811463046757, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4516126', 'n_tokens_mistral': 983, 'n_tokens_neox': 949, 'n_words': 200}