<commit_before>import tensorflow as tf
import numpy as np
from time import time
import os
from six.moves.cPickle import dump, load

from .classifier import Classifier
from .utils import reshape_marginals, LabelBalancer

class TFNoiseAwareModel(Classifier):
    """
    Generic NoiseAwareModel class for TensorFlow models.
    Note that the actual network is built when train is called (to allow for
    model architectures which depend on the training data, e.g. vocab size).
    
    :param n_threads: Parallelism to use; single-threaded if None
    :param seed: Top level seed which is passed into both numpy operations
        via a RandomState maintained by the class, and into TF as a graph-level
        seed.
    """
    def __init__(self, n_threads=None, seed=123, **kwargs):
        self.n_threads = n_threads
        self.seed = seed
        self.rand_state = np.random.RandomState()
        super(TFNoiseAwareModel, self).__init__(**kwargs)

    def _build_model(self, **model_kwargs):
        """
        Builds the TensorFlow Operations for the model. Must set the following:
            - self.logits: The un-normalized potentials for the variables
                ("logits" in keeping with TensorFlow terminology)
            - self.Y: The training marginals to fit to
            - self.marginals_op: Normalized predicted marginals for the vars

        Additional ops must be set depending on whether the default
        self._construct_feed_dict method below is used, or a custom one.

        Note that _build_model is called in the train method, allowing for the 
        network to be constructed dynamically depending on the training set 
        (e.g., the size of the vocabulary for a text LSTM model)

        Note also that model_kwargs are saved to disk by the self.save method,
        as they are needed to rebuild / reload the model. *All hyperparameters
        needed to rebuild the model must be passed in here for model reloading
        to work!*
        """
        raise NotImplementedError()

    def _build_training_ops(self, **training_kwargs):
        """
        Builds the TensorFlow Operations for the training procedure. Must set 
        the following:
            - self.loss: Loss function
            - self.optimizer: Training operation
        """
        # Define loss and marginals ops
        if self.cardinality > 2:
            loss_fn = tf.nn.softmax_cross_entropy_with_logits
        else:
            loss_fn = tf.nn.sigmoid_cross_entropy_with_logits
        self.loss = tf.reduce_mean(loss_fn(logits=self.logits, labels=self.Y))
        
        # Build training op
        self.lr = tf.placeholder(tf.float32)
        self.optimizer = tf.train.AdamOptimizer(self.lr).minimize(self.loss)

    def _construct_feed_dict(self, X_b, Y_b, lr=0.01, **kwargs):
        """
        Given a batch of data and labels, and other necessary hyperparams,
        construct a python dictionary to use in the training step as feed_dict.
<commit_msg>Removing use of reduce_mean in disc_learning to test non-determinism on GPU<commit_after>import tensorflow as tf
import numpy as np
from time import time
import os
from six.moves.cPickle import dump, load

from .classifier import Classifier
from .utils import reshape_marginals, LabelBalancer

class TFNoiseAwareModel(Classifier):
    """
    Generic NoiseAwareModel class for TensorFlow models.
    Note that the actual network is built when train is called (to allow for
    model architectures which depend on the training data, e.g. vocab size).
    
    :param n_threads: Parallelism to use; single-threaded if None
    :param seed: Top level seed which is passed into both numpy operations
        via a RandomState maintained by the class, and into TF as a graph-level
        seed.
    """
    def __init__(self, n_threads=None, seed=123, **kwargs):
        self.n_threads = n_threads
        self.seed = seed
        self.rand_state = np.random.RandomState()
        super(TFNoiseAwareModel, self).__init__(**kwargs)

    def _build_model(self, **model_kwargs):
        """
        Builds the TensorFlow Operations for the model. Must set the following:
            - self.logits: The un-normalized potentials for the variables
                ("logits" in keeping with TensorFlow terminology)
            - self.Y: The training marginals to fit to
            - self.marginals_op: Normalized predicted marginals for the vars

        Additional ops must be set depending on whether the default
        self._construct_feed_dict method below is used, or a custom one.

        Note that _build_model is called in the train method, allowing for the 
        network to be constructed dynamically depending on the training set 
        (e.g., the size of the vocabulary for a text LSTM model)

        Note also that model_kwargs are saved to disk by the self.save method,
        as they are needed to rebuild / reload the model. *All hyperparameters
        needed to rebuild the model must be passed in here for model reloading
        to work!*
        """
        raise NotImplementedError()

    def _build_training_ops(self, **training_kwargs):
        """
        Builds the TensorFlow Operations for the training procedure. Must set 
        the following:
            - self.loss: Loss function
            - self.optimizer: Training operation
        """
        # Define loss and marginals ops
        if self.cardinality > 2:
            loss_fn = tf.nn.softmax_cross_entropy_with_logits
        else:
            loss_fn = tf.nn.sigmoid_cross_entropy_with_logits
       
        # self.loss = tf.reduce_mean(loss_fn(logits=self.logits, labels=self.Y))

        # Avoiding the use of non-deterministic reduce_ ops
        l = tf.reshape(loss_fn(logits=self.logits, labels=self.Y), [1, -1])
        self.loss = tf.squeeze(tf.matmul(l, tf.ones_like(l), transpose_b=True))
        self.loss /= tf.cast(tf.shape(l)[1], tf.float32)
        
        # Build training op
        self.lr = tf.placeholder(tf.float32)
        self.optimizer = tf.train.AdamOptimizer(self.lr).minimize(self.loss)

    def _construct_feed_dict(self, X_b, Y_b, lr=0.01, **kwargs):
        """
        Given a batch of data and labels, and other necessary hyperparams,
        construct a python dictionary to use in the training step as feed_dict.

{'fraction_non_alphanumeric': 0.06617182564750473, 'fraction_numerical': 0.0036323436512950094, 'mean_word_length': 2.6250715512306813, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '4312215', 'n_tokens_mistral': 1766, 'n_tokens_neox': 1709, 'n_words': 657}