def linear_mapping(x, y, n_jobs=8):
  # train least square
  projection = linear_model.LinearRegression(n_jobs=n_jobs)
  projection.fit(x, y)
  return projection.predict def gan_mapping(x, y, lr=1e-4, lmbda=10., gamma=0., beta1=0.5,
                activation=tf.nn.relu, epoch=30, disc_iters=10,
                batch_size=128):
  n_data, x_dim = x.shape
  y_dim = y.shape[1]
  model = WGANGP(x_dim, y_dim, lr=lr, lmbda=lmbda, gamma=gamma, beta1=beta1,
                 activation=activation)

  gen_sampler = inf_batch_iterator(n_data, batch_size)
  disc_sampler = inf_batch_iterator(n_data, batch_size)
  num_batch_per_epoch = n_data // batch_size + (n_data % batch_size) != 0

  sess = tf.Session()
  sess.run(tf.global_variables_initializer())
  for e in range(epoch):
    train_d_loss = []
    train_g_loss = []
    train_l2_loss = []
    for _ in range(num_batch_per_epoch):
      # train disc first
      for _ in range(disc_iters):
        disc_idx = next(disc_sampler)
        disc_x, disc_y = x[disc_idx], y[disc_idx]
        d_err = model.train_disc_one_batch(sess, disc_x, disc_y)
        train_d_loss.append(d_err)

      gen_idx = next(gen_sampler)
      gen_x, gen_y = x[gen_idx], y[gen_idx]
      g_err, l2_err = model.train_gen_one_batch(sess, gen_x, gen_y)
      train_g_loss.append(g_err)
      train_l2_loss.append(l2_err)

    train_d_loss = np.mean(train_d_loss)
    train_g_loss = np.mean(train_g_loss)
    train_l2_loss = np.mean(train_l2_loss)
    log('Epoch: {}, disc loss: {:.4f}, gen loss: {:.4f},'
        ' l2 loss: {:.4f}'.format(
          e + 1, train_d_loss, train_g_loss, train_l2_loss))

  def mapping(z):
    mapped = []
    for idx in iterate_minibatches_indices(len(z), batch_size=2048):
      batch_mapped = model.generate(sess, z[idx])
      mapped.append(batch_mapped)

    tf.keras.backend.clear_session()
    return np.vstack(mapped)

  return mapping def mlp_mapping(x, y, epochs=30, batch_size=256, lr=1e-3, wd=1e-4,
                activation=tf.nn.relu, norm=False):
  dim = y.shape[1]
  l2_reg = tf.keras.regularizers.l2(wd)
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.Dense(dim, activation=activation, name='reg_fc',
                                  kernel_regularizer=l2_reg))
  model.add(tf.keras.layers.Dense(dim, name='reg_output',
                                  kernel_regularizer=l2_reg))
  if norm:
    model.add(tf.keras.layers.Lambda(lambda x: tf.nn.l2_normalize(x, axis=-1),
                                     name='reg_norm'))
  model.compile(tf.keras.optimizers.Adam(lr), loss='mse', metric='mse')
  model.fit(x, y, batch_size=batch_size, verbose=1, epochs=epochs)

  def mapping(z):
    mapped = model.predict(z, batch_size=1024)
    tf.keras.backend.clear_session()
    return mapped

  return mapping def load_author_data(min_book=2, train_size=50, test_size=300,
                     unlabeled_size=0, top_attr=0, split_by_book=True,
                     split_word=True, remove_punct=False, seed=54321,
                     min_len=0):

  min_size = max(test_size + unlabeled_size, test_size * 2)
  if split_by_book:
    # use sentence from one book to infer authorship of other books
    assert min_book >= 2

  filenames = os.listdir(BOOKCORPUS_AUTHOR_DIR)
  attribute_dict = load_book_author(filenames)
  author_book_count = Counter(attribute_dict.values())
  filtered_filenames = [fname for fname in filenames if
                        author_book_count[attribute_dict[fname]] >= min_book]

  if remove_punct:
    file_sents = Parallel(n_jobs=8)(
      delayed(load_single_author_file)(fname, split_word, remove_punct, min_len)
      for fname in tqdm.tqdm(filtered_filenames)
    )
  else:
    file_sents = [load_single_author_file(fname, split_word,
                                          remove_punct, min_len)
                  for fname in tqdm.tqdm(filtered_filenames)]

  author_file_sents = defaultdict(list)
  author_sent_counts = Counter()

  for fname, sents in zip(filtered_filenames, file_sents):
    author = attribute_dict[fname]
    author_file_sents[author].append(sents)
    author_sent_counts[author] += len(sents)

  authors = [tup[0] for tup in author_sent_counts.most_common()]
  filtered_authors = []
  for author in authors:
    author_books = author_file_sents[author]
    num_books = len(author_books)
    assert num_books >= min_book
    if all(len(sents) >= min_size for sents in author_books):
      filtered_authors.append(author)
      if 0 < top_attr <= len(filtered_authors):
        break

  train_sents, train_authors = [], []
  test_sents, test_authors = [], []
  unlabeled_sents, unlabeled_authors = [], []

  np.random.seed(seed)
  for author in filtered_authors:
    author_books = author_file_sents[author]
    num_books = len(author_books)

    if split_by_book:
      train_author_book_idx = np.random.choice(num_books)
      train_author_book = author_books[train_author_book_idx]
      train_author_sents = np.random.choice(train_author_book,
                                            size=train_size, replace=False)
      rest_author_book_indices = [i for i in range(num_books)
                                  if i != train_author_book_idx]
      rest_author_books = np.concatenate([author_books[i] for i in
                                          rest_author_book_indices])
      rest_author_sents = np.random.choice(rest_author_books,
                                           size=min_size, replace=False)
    else:
      all_author_sents = np.concatenate(author_books)
      sampled_author_sents = np.random.choice(all_author_sents,
                                              size=train_size + min_size,
                                              replace=False)
      rest_author_sents = sampled_author_sents[:min_size]
      train_author_sents = sampled_author_sents[min_size: min_size + train_size]

    test_author_sents = rest_author_sents[:test_size]
    if unlabeled_size > 0:
      unlabeled_author_sents = rest_author_sents[test_size:
                                                 test_size + unlabeled_size]
      unlabeled_sents.append(unlabeled_author_sents)
      unlabeled_authors.extend([author] * len(unlabeled_author_sents))

    train_sents.append(train_author_sents)
    train_authors.extend([author] * len(train_author_sents))
    test_sents.append(test_author_sents)
    test_authors.extend([author] * len(test_author_sents))

  train_sents = np.concatenate(train_sents)
  test_sents = np.concatenate(test_sents)

  if len(unlabeled_sents) > 0:
    unlabeled_sents = np.concatenate(unlabeled_sents)

  return train_sents, train_authors, test_sents, test_authors, \
      unlabeled_sents, unlabeled_authors def iterate_minibatches_indices(n, batch_size, shuffle=False,
                                include_last=True):
  indices = np.arange(n)

  num_batches = n // batch_size
  if include_last:
    num_batches += (n % batch_size != 0)

  if num_batches == 0:
    num_batches = 1

  if shuffle:
    np.random.shuffle(indices)

  for batch_idx in range(num_batches):
    batch_indices = indices[batch_idx * batch_size:
                            (batch_idx + 1) * batch_size]
    yield batch_indices def get_similarity_metric(x, y, metric, rtn_loss=False):
  if metric == 'cosine':
    x = tf.nn.l2_normalize(x, axis=-1)
    y = tf.nn.l2_normalize(y, axis=-1)
    sim = tf.reduce_sum(tf.multiply(x, y), axis=1)
    loss = 1 - sim
  elif metric == 'log':
    dot = tf.reduce_sum(tf.multiply(x, y), axis=1)
    sim = tf.log_sigmoid(dot)
    loss = -sim
  elif metric == 'dot':
    sim = tf.reduce_sum(tf.multiply(x, y), axis=1)
    loss = -sim
  elif metric == 'l2':
    # sim = -tf.reduce_sum(tf.square(x - y), axis=-1)
    sim = -tf.sqrt(tf.reduce_sum(tf.square(x - y), axis=-1))
    loss = -sim   # / tf.reduce_sum(tf.square(y), axis=-1, keepdims=True)
  else:
    raise ValueError('No such metric', metric)
  return loss if rtn_loss else sim def log(msg):
  if msg[-1] != '\n':
    msg += '\n'
  sys.stderr.write(msg)
  sys.stderr.flush() def sents_to_labels(sents, label_thresh=0):
  labels = []
  for s in sents:
    label = s[s > label_thresh]
    assert len(label) > 0, s
    labels.append(label)
  return np.asarray(labels) def tp_fp_fn_metrics_np(y_pred, y_true, rtn_match=False, rare_thresh=0):
  tp, fp, fn = 0., 0., 0.
  matched = {}
  i = 0
  for p, t in zip(y_pred, y_true):
    t = np.unique(t[t > 0])
    p = np.unique(p[p > 0])
    matched_all = np.intersect1d(t, p)

    if rtn_match:
      matched_rare = np.intersect1d(t[t > rare_thresh], p[p > rare_thresh])
      if len(matched_rare) >= 1:
        matched[i] = matched_all

    tp += len(matched_all)  # predicted and in label
    fp += len(np.setdiff1d(p, t))  # predicted but not in label
    fn += len(np.setdiff1d(t, p))  # in label but not predicted
    i += 1

  if rtn_match:
    return tp, fp, fn, matched

  return tp, fp, fn def sinkhorn(log_alpha, n_iters=5):
  n = tf.shape(log_alpha)[1]
  log_alpha = tf.reshape(log_alpha, [-1, n, n])

  for _ in range(n_iters):
    log_alpha -= tf.reshape(tf.reduce_logsumexp(log_alpha, axis=2), [-1, n, 1])
    log_alpha -= tf.reshape(tf.reduce_logsumexp(log_alpha, axis=1), [-1, 1, n])
  return tf.exp(log_alpha)
from absl import app
from absl import flags
from train_feature_mapper import linear_mapping, gan_mapping, mlp_mapping
from data.bookcorpus import load_author_data as load_bookcorpus_author
from utils.sent_utils import iterate_minibatches_indices, get_similarity_metric
from utils.common_utils import log
from invert.utils import sents_to_labels, tp_fp_fn_metrics_np, sinkhorn
import time
import tqdm
import tensorflow as tf
import tensorflow_hub as hub
import sentencepiece as spm
import numpy as np
flags.DEFINE_integer(
    "max_seq_length", 32,
    "The maximum total input sequence length after WordPiece tokenization. "
    "Sequences longer than this will be truncated, and sequences shorter "
    "than this will be padded.")
flags.DEFINE_string('metric', 'l2', 'Metric to optimize')
flags.DEFINE_string('mapper', 'linear', 'Mapper to use')
flags.DEFINE_string('model', 'multiset', 'Model for learning based inversion')
flags.DEFINE_boolean(
  'learning', False,
  'Learning based inversion or optimize based')
flags.DEFINE_boolean(
  'cross_domain', False,
  'Cross domain data for learning based inversion')
flags.DEFINE_float('temp', 0.1, 'Temperature for optimization')
flags.DEFINE_float('lr', 1e-3, 'Learning rate')
flags.DEFINE_float('alpha', 0.0, 'Coefficient for regularization')
flags.DEFINE_float('C', 0.0, 'Label distribution aware margin')
flags.DEFINE_float('wd', 1e-4, 'Weight decay')

FLAGS = flags.FLAGS


LAYER_NAMES = [
  "module_apply_default/Encoder_en/KonaTransformer/Encode/Scatter/ScatterNd:0",
  "module_apply_default/Encoder_en/KonaTransformer/Encode/TransformerStack/"
  "Layer_0/TransformerLayer/FFN/layer_postprocess/add:0",
  "module_apply_default/Encoder_en/KonaTransformer/Encode/TransformerStack/"
  "Layer_1/TransformerLayer/FFN/layer_postprocess/add:0",
  "module_apply_default/Encoder_en/hidden_layers/l2_normalize:0",
]


def sents_to_sparse(ids):
  max_len = max(len(x) for x in ids)
  dense_shape = (len(ids), max_len)
  values = [item for sublist in ids for item in sublist]
  indices = [[row, col] for row in range(len(ids)) for col in
             range(len(ids[row]))]

  return values, indices, dense_shape


def mean_pool(tensor):
  denom = tf.get_default_graph().get_tensor_by_name(
    "module_apply_default/Encoder_en/KonaTransformer/ExpandDims_1:0")
  masks = tf.get_default_graph().get_tensor_by_name(
    "module_apply_default/Encoder_en/KonaTransformer/ExpandDims:0")
  return tf.reduce_sum(tensor * masks, axis=1) / denom


def get_fetch_by_layer(layer_idx):
  fetch = tf.get_default_graph().get_tensor_by_name(LAYER_NAMES[layer_idx])
  if fetch.get_shape().ndims == 3:
    fetch = mean_pool(fetch)
  assert fetch.get_shape().ndims == 2
  return fetch


def load_inversion_data():
  module = hub.Module(
    "https://tfhub.dev/google/universal-sentence-encoder-lite/2")

  sp = spm.SentencePieceProcessor()
  sp.Load(SPM_MODEL_PATH)

  input_placeholder = tf.sparse_placeholder(tf.int64,
                                            shape=[None, None],
                                            name='sparse_placeholder')
  module(inputs=dict(
      values=input_placeholder.values,
      indices=input_placeholder.indices,
      dense_shape=input_placeholder.dense_shape))

  learn_mapping = FLAGS.high_layer_idx != FLAGS.low_layer_idx
  if learn_mapping:
    outputs = [get_fetch_by_layer(FLAGS.low_layer_idx),
               get_fetch_by_layer(FLAGS.high_layer_idx)]
  else:
    outputs = get_fetch_by_layer(FLAGS.low_layer_idx)

  train_sents, _, test_sents, _, _, _ = load_bookcorpus_author(
      train_size=FLAGS.train_size, test_size=FLAGS.test_size,
      unlabeled_size=0, split_by_book=True, split_word=False,
      top_attr=800, remove_punct=False)

  sess = tf.Session()
  sess.run(tf.global_variables_initializer())

  def encode_sents(sents):
    y = [np.asarray(sp.EncodeAsIds(x)[:FLAGS.max_seq_length])
         for x in train_sents]
    y = np.asarray(y)
    n_data = len(sents)
    pbar = tqdm.tqdm(total=n_data)
    embs_low, embs_high = [], []
    for b_idx in iterate_minibatches_indices(n_data, 512):
      values, indices, dense_shape = sents_to_sparse(y[b_idx])
      emb = sess.run(outputs,
                     feed_dict={input_placeholder.values: values,
                                input_placeholder.indices: indices,
                                input_placeholder.dense_shape: dense_shape})
      if learn_mapping:
        embs_low.append(emb[0])
        embs_high.append(emb[1])
      else:
        embs_low.append(emb)
      pbar.update(len(b_idx))

    pbar.close()
    if learn_mapping:
      return [np.vstack(embs_low), np.vstack(embs_high)], y
    else:
      return np.vstack(embs_low), y

  train_x, train_y = encode_sents(train_sents)
  test_x, test_y = encode_sents(test_sents)
  tf.keras.backend.clear_session()

  if learn_mapping:

    log('Training high to low mapping...')