get_language_model: just return hclg

Author: maxhawkins

gentle/language_model.py

@@ -50,7 +50,10 @@ def get_node_id(word):
     return output
 
 def get_language_model(kaldi_seq, proto_langdir='PROTO_LANGDIR'):
-    """Generates a language model to fit the text
+    """Generates a language model to fit the text.
+
+    Returns the filename of the generated language model FST.
+    The caller is resposible for removing the generated file.
 
     `proto_langdir` is a path to a directory containing prototype model data
     `kaldi_seq` is a list of words within kaldi's vocabulary.
@@ -62,22 +65,21 @@ def get_language_model(kaldi_seq, proto_langdir='PROTO_LANGDIR'):
     txt_fst_file.write(txt_fst)
     txt_fst_file.close()
 
-    out_dir = tempfile.mkdtemp()
-
+    hclg_filename = tempfile.mktemp(suffix='_HCLG.fst')
     try:
         subprocess.check_output([MKGRAPH_PATH,
                          os.path.join(proto_langdir, 'langdir'),
                          os.path.join(proto_langdir, 'modeldir'),
                          txt_fst_file.name,
                          os.path.join(proto_langdir, "graphdir/words.txt"),
-                         os.path.join(out_dir, 'HCLG.fst')])
+                         hclg_filename])
     except Exception, e:
-        shutil.rmtree(out_dir)
+        os.unlink(hclg_filename)
         raise e
     finally:
         os.unlink(txt_fst_file.name)
 
-    return out_dir
+    return hclg_filename
 
 if __name__=='__main__':
     import sys

gentle/language_model_transcribe.py

@@ -27,18 +27,17 @@ def lm_transcribe(audio_f, transcript, proto_langdir, nnet_dir,
 
     ks = ms.get_kaldi_sequence()
 
-    gen_model_dir = language_model.get_language_model(ks, proto_langdir)
+    gen_hclg_filename = language_model.get_language_model(ks, proto_langdir)
     try:
-        gen_hclg_path = os.path.join(gen_model_dir, 'HCLG.fst')
-        k = standard_kaldi.Kaldi(nnet_dir, gen_hclg_path, proto_langdir)
+        k = standard_kaldi.Kaldi(nnet_dir, gen_hclg_filename, proto_langdir)
 
         trans = standard_kaldi.transcribe(k, audio_f,
                                           partial_results_cb=partial_cb,
                                           partial_results_kwargs=partial_kwargs)
 
         ret = diff_align.align(trans["words"], ms)
     finally:
-        shutil.rmtree(gen_model_dir)
+        os.unlink(gen_hclg_filename)
 
     return {
         "transcript": transcript,