problem_name = "librispeech_clean"
asr_problem = problems.problem(problem_name)
encoders = asr_problem.feature_encoders(None)

model_name = "transformer"
hparams_set = "transformer_librispeech_tpu"

hparams = trainer_lib.create_hparams(hparams_set,data_dir=data_dir, problem_name=problem_name)
asr_model = registry.model(model_name)(hparams, Modes.PREDICT)

def encode(x):
  waveforms = encoders["waveforms"].encode(x)
  encoded_dict = asr_problem.preprocess_example({"waveforms":waveforms, "targets":[]}, Modes.PREDICT, hparams)
  
  return {"inputs" : tf.expand_dims(encoded_dict["inputs"], 0), "targets" : tf.expand_dims(encoded_dict["targets"], 0)}

def decode(integers):
  integers = list(np.squeeze(integers))
  if 1 in integers:
    integers = integers[:integers.index(1)]
  return encoders["targets"].decode(np.squeeze(integers))


"""
Explanation: Define problem, hparams, model, encoder and decoder
Definition of this model (as well as many more) can be found on tensor2tensor github page.
End of explanation
"""


# Copy the pretrained checkpoint locally
ckpt_name = "transformer_asr_180214"
gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)
print(gs_ckpt)
!gsutil cp -R {gs_ckpt} {checkpoint_dir} 
ckpt_path = tf.train.latest_checkpoint(os.path.join(checkpoint_dir, ckpt_name))
ckpt_path

"""
Explanation: Define path to checkpoint
In this demo we are using a pretrained model.
Instructions for training your own model can be found in the tutorial on tensor2tensor page.
End of explanation
"""


# Restore and transcribe!
def transcribe(inputs):
  encoded_inputs = encode(inputs)
  with tfe.restore_variables_on_create(ckpt_path): 
    model_output = asr_model.infer(encoded_inputs, beam_size=2, alpha=0.6, decode_length=1)["outputs"]
  return decode(model_output)

def play_and_transcribe(inputs):
  waveforms = encoders["waveforms"].encode(inputs)
  IPython.display.display(IPython.display.Audio(data=waveforms, rate=16000))
  return transcribe(inputs) 

"""
Explanation: Define transcribe function
End of explanation
"""


uploaded = google.colab.files.upload()
prerecorded_messages = []

for fn in uploaded.keys():
  print('User uploaded file "{name}" with length {length} bytes'.format(
      name=fn, length=len(uploaded[fn])))
  mem_file = cStringIO.StringIO(uploaded[fn])
    
  save_filename = os.path.join(tmp_dir, fn)
  with open(save_filename, 'w') as fd:
    mem_file.seek(0)
    shutil.copyfileobj(mem_file, fd)
  prerecorded_messages.append(save_filename)
         
              
for inputs in prerecorded_messages:
  outputs = play_and_transcribe(inputs)

  print("Inputs: %s" % inputs)
  print("Outputs: %s" % outputs)

"""
Explanation: Decoding prerecorded examples
You can upload any .wav files. They will be transcribed if frame rate matches Librispeeche's frame rate (16kHz).
End of explanation
"""


# Records webm file and converts
def RecordNewAudioSample(filename=None, webm_filename=None):
  """Args:
          filename - string, path for storing wav file
          webm_filename - string, path for storing webm file
     Returns:
          string - path where wav file was saved. (=filename if specified)
    
  """
  # Create default filenames in tmp_dir if not specified.
  if not filename:
    filename = os.path.join(tmp_dir, "recording.wav")
  if not webm_filename:
    webm_filename = os.path.join(tmp_dir, "recording.webm")
    
  # Record webm file form colab.
  
  audio = google.colab._message.blocking_request('user_media', {"audio":True, "video":False, "duration":-1}, timeout_sec=600)
  #audio = frontend.RecordMedia(True, False)
  
  # Convert the recording into in_memory file.
  music_mem_file = cStringIO.StringIO(
      base64.decodestring(audio[audio.index(',')+1:]))
  
  # Store webm recording in webm_filename. Storing is necessary for conversion.
  with open(webm_filename, 'w') as fd:
    music_mem_file.seek(0)
    shutil.copyfileobj(music_mem_file, fd)
    
  # Open stored file and save it as wav with sample_rate=16000.
  pydub.AudioSegment.from_file(webm_filename, codec="opus"
                              ).set_frame_rate(16000).export(out_f=filename,
                                                             format="wav")
  return filename

# Record the sample
my_sample_filename = RecordNewAudioSample()
print my_sample_filename

print play_and_transcribe(my_sample_filename)

"""
Explanation: Recording your own examples
End of explanation
"""


