# Path: px/nmt/environment_client.py
# def multi_call_environment(pool, stub, request, timeouts):
# def single_call_environment(args):
#   def __init__(self, *args, **kwargs):
#   def __setitem__(self, key, value):
#   def _check_size_limit(self):
# def make_cache_key(question, docid):
# def make_environment_reward_fn(environment_server,
#                                timeouts=None,
#                                mode='squad',
#                                use_cache=False,
#                                cache_size=-1,
#                                env_call_parallelism=1):
#   def environment_reward_fn(questions, doc_ids):
# class LimitedSizeDict(OrderedDict):
# 
# Path: px/nmt/reformulator.py
# def load_hparams(hparams_file, out_dir):
#   def __init__(self, hparams_path, source_prefix, out_dir,
#                environment_server_address):
#   def tokenize(self, questions, prefix=""):
#   def detokenize(self, text):
#   def reformulate(self, questions, inference_mode, trie_excludes=None):
#   def train(self, sources, annotations):
# class Reformulator(object):
# 
# Path: px/nmt/utils/logging_utils.py
# def safe_string(s):
#   """Safely converts unicode and plain strings to byte strings."""
#   if isinstance(s, unicode):
#     try:
#       s = s.encode('utf-8')
#     except UnicodeDecodeError:
#       s = repr(s)
#   return s
# 
# Path: px/selector/selector_keras.py
# FLAGS = flags.FLAGS
# class Selector(object):
#   def __init__(self):
#   def load(self, name):
#   def save(self, name):
#   def _get_checkpoint_paths(self, name):
#   def _build_embedding_matrix(self):
#   def _build_model(self, embedding_matrix):
#   def encode_labels(self, labels):
#   def encode_texts(self, texts):
#   def encode_data(self, questions, original_questions, answers, labels):
#   def encode_train(self, question_lists, answer_lists, score_lists):
#   def train(self, questions, answers, scores):
#   def eval(self, question_lists, answer_lists, score_lists):
# 
# Path: px/utils/eval_utils.py
# def compute_f1_single(prediction, ground_truth):
# def compute_f1(prediction, gold_answers):
# 
# Path: third_party/nmt/utils/misc_utils.py
# def check_tensorflow_version():
# def safe_exp(value):
# def print_time(s, start_time):
# def print_out(s, f=None, new_line=True):
# def print_hparams(hparams, skip_patterns=None, header=None):
# def load_hparams(model_dir):
# def maybe_parse_standard_hparams(hparams, hparams_path):
# def save_hparams(out_dir, hparams):
# def debug_tensor(s, msg=None, summarize=10):
# def add_summary(summary_writer, global_step, tag, value):
# def get_config_proto(log_device_placement=False, allow_soft_placement=True,
#                      num_intra_threads=0, num_inter_threads=0):
# def format_text(words):
# def format_bpe_text(symbols, delimiter=b"@@"):
# def format_spm_text(symbols):


import codecs
import json
import os
import random
import time
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from px.nmt import environment_client
from px.nmt import reformulator
from px.nmt.utils.logging_utils import safe_string
from px.proto import reformulator_pb2
from px.selector import selector_keras as selector
from px.utils import eval_utils
from third_party.nmt.utils import misc_utils
      FLAGS.environment_server_address,
      mode=FLAGS.mode,
      env_call_parallelism=FLAGS.env_sample_parallelism)

  eval_environment_fn = environment_client.make_environment_reward_fn(
      FLAGS.environment_server_address,
      mode='searchqa',
      env_call_parallelism=FLAGS.env_eval_parallelism)

  # Read data.
  questions, annotations, docid_2_answer = read_data(
      questions_file=FLAGS.train_questions,
      annotations_file=FLAGS.train_annotations,
      answers_file=FLAGS.train_data,
      preprocessing_mode=FLAGS.mode)
  dev_questions, dev_annotations, dev_docid_2_answer = read_data(
      questions_file=FLAGS.dev_questions,
      annotations_file=FLAGS.dev_annotations,
      answers_file=FLAGS.dev_data,
      preprocessing_mode=FLAGS.mode,
      max_lines=FLAGS.max_dev_examples)

  # Summary writer that writes events to a folder. TensorBoard will later read
  # from it.
  summary_writer = tf.summary.FileWriter(
      os.path.join(
          FLAGS.tensorboard_dir,
          'reformulator_and_selector_training_log_' + str(time.time())))

  if FLAGS.enable_selector_training:

# Path: px/nmt/environment_client.py
# def multi_call_environment(pool, stub, request, timeouts):
# def single_call_environment(args):
#   def __init__(self, *args, **kwargs):
#   def __setitem__(self, key, value):
#   def _check_size_limit(self):
# def make_cache_key(question, docid):
# def make_environment_reward_fn(environment_server,
#                                timeouts=None,
#                                mode='squad',
#                                use_cache=False,
#                                cache_size=-1,
#                                env_call_parallelism=1):
#   def environment_reward_fn(questions, doc_ids):
# class LimitedSizeDict(OrderedDict):
# 
# Path: px/nmt/reformulator.py
# def load_hparams(hparams_file, out_dir):
#   def __init__(self, hparams_path, source_prefix, out_dir,
#                environment_server_address):
#   def tokenize(self, questions, prefix=""):
#   def detokenize(self, text):
#   def reformulate(self, questions, inference_mode, trie_excludes=None):
#   def train(self, sources, annotations):
# class Reformulator(object):
# 
# Path: px/nmt/utils/logging_utils.py
# def safe_string(s):
#   """Safely converts unicode and plain strings to byte strings."""
#   if isinstance(s, unicode):
#     try:
#       s = s.encode('utf-8')
#     except UnicodeDecodeError:
#       s = repr(s)
#   return s
# 
# Path: px/selector/selector_keras.py
# FLAGS = flags.FLAGS
# class Selector(object):
#   def __init__(self):
#   def load(self, name):
#   def save(self, name):
#   def _get_checkpoint_paths(self, name):
#   def _build_embedding_matrix(self):
#   def _build_model(self, embedding_matrix):
#   def encode_labels(self, labels):
#   def encode_texts(self, texts):
#   def encode_data(self, questions, original_questions, answers, labels):
#   def encode_train(self, question_lists, answer_lists, score_lists):
#   def train(self, questions, answers, scores):
#   def eval(self, question_lists, answer_lists, score_lists):
# 
# Path: px/utils/eval_utils.py
# def compute_f1_single(prediction, ground_truth):
# def compute_f1(prediction, gold_answers):
# 
# Path: third_party/nmt/utils/misc_utils.py
# def check_tensorflow_version():
# def safe_exp(value):
# def print_time(s, start_time):
# def print_out(s, f=None, new_line=True):
# def print_hparams(hparams, skip_patterns=None, header=None):
# def load_hparams(model_dir):
# def maybe_parse_standard_hparams(hparams, hparams_path):
# def save_hparams(out_dir, hparams):
# def debug_tensor(s, msg=None, summarize=10):
# def add_summary(summary_writer, global_step, tag, value):
# def get_config_proto(log_device_placement=False, allow_soft_placement=True,
#                      num_intra_threads=0, num_inter_threads=0):
# def format_text(words):
# def format_bpe_text(symbols, delimiter=b"@@"):
# def format_spm_text(symbols):


# Path: px/nmt/reformulator_and_selector_training.py
import codecs
import json
import os
import random
import time
import numpy as np
import tensorflow as tf
from absl import app
from absl import flags
from px.nmt import environment_client
from px.nmt import reformulator
from px.nmt.utils.logging_utils import safe_string
from px.proto import reformulator_pb2
from px.selector import selector_keras as selector
from px.utils import eval_utils
from third_party.nmt.utils import misc_utils

      FLAGS.environment_server_address,
      mode=FLAGS.mode,
      env_call_parallelism=FLAGS.env_sample_parallelism)

  eval_environment_fn = environment_client.make_environment_reward_fn(
      FLAGS.environment_server_address,
      mode='searchqa',
      env_call_parallelism=FLAGS.env_eval_parallelism)

  # Read data.
  questions, annotations, docid_2_answer = read_data(
      questions_file=FLAGS.train_questions,
      annotations_file=FLAGS.train_annotations,
      answers_file=FLAGS.train_data,
      preprocessing_mode=FLAGS.mode)
  dev_questions, dev_annotations, dev_docid_2_answer = read_data(
      questions_file=FLAGS.dev_questions,
      annotations_file=FLAGS.dev_annotations,
      answers_file=FLAGS.dev_data,
      preprocessing_mode=FLAGS.mode,
      max_lines=FLAGS.max_dev_examples)

  # Summary writer that writes events to a folder. TensorBoard will later read
  # from it.
  summary_writer = tf.summary.FileWriter(
      os.path.join(
          FLAGS.tensorboard_dir,
          'reformulator_and_selector_training_log_' + str(time.time())))

  if FLAGS.enable_selector_training:

    selector_model = selector.Selector()