# 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
        responses_beam = reformulator_instance.reformulate(
            questions=questions_batch,
            inference_mode=reformulator_pb2.ReformulatorRequest.BEAM_SEARCH)

        # Discard answers.
        reformulations_beam = [
            [rf.reformulation for rf in rsp] for rsp in responses_beam
        ]

      if FLAGS.enable_reformulator_training:
        # Train reformulator model.
        if FLAGS.debug:
          print('Training reformulator...')

        reformulator_loss, f1s, reformulations = reformulator_instance.train(
            sources=questions_batch, annotations=annotations_batch)

        f1_avg = f1s.mean()

        if [] in reformulations:
          if FLAGS.debug:
            print('Found empty rewrites! Skipping this batch.')
          continue

        if FLAGS.debug:
          print('Rewrite: {}'.format(safe_string(reformulations[0])))
          print('Avg F1: {}'.format(f1_avg))
          print('Loss  : {}'.format(reformulator_loss))

        # Write the f1_avg and loss to Tensorboard.

# 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

        responses_beam = reformulator_instance.reformulate(
            questions=questions_batch,
            inference_mode=reformulator_pb2.ReformulatorRequest.BEAM_SEARCH)

        # Discard answers.
        reformulations_beam = [
            [rf.reformulation for rf in rsp] for rsp in responses_beam
        ]

      if FLAGS.enable_reformulator_training:
        # Train reformulator model.
        if FLAGS.debug:
          print('Training reformulator...')

        reformulator_loss, f1s, reformulations = reformulator_instance.train(
            sources=questions_batch, annotations=annotations_batch)

        f1_avg = f1s.mean()

        if [] in reformulations:
          if FLAGS.debug:
            print('Found empty rewrites! Skipping this batch.')
          continue

        if FLAGS.debug:
          print('Rewrite: {}'.format(safe_string(reformulations[0])))
          print('Avg F1: {}'.format(f1_avg))
          print('Loss  : {}'.format(reformulator_loss))

        # Write the f1_avg and loss to Tensorboard.

        misc_utils.add_summary(