MMI support

1 files changed, 59 insertions, 1 deletions

diff --git a/kaldi_seq/src/init.c b/kaldi_seq/src/init.c
index 88d0a80..9b38056 100644
--- a/kaldi_seq/src/init.c
+++ b/kaldi_seq/src/init.c
@@ -1,8 +1,10 @@
 #include "nerv/common.h"
 #include "kaldi_mpe.h"
+#include "kaldi_mmi.h"
 #include <stdio.h>
 
 const char *nerv_kaldi_mpe_tname = "nerv.KaldiMPE";
+const char *nerv_kaldi_mmi_tname = "nerv.KaldiMMI";
 const char *nerv_matrix_cuda_float_tname = "nerv.CuMatrixFloat";
 const char *nerv_matrix_host_float_tname = "nerv.MMatrixFloat";
 
@@ -63,11 +65,67 @@ static const luaL_Reg mpe_methods[] = {
 
 static void mpe_init(lua_State *L) {
     luaT_newmetatable(L, nerv_kaldi_mpe_tname, NULL,
-                        mpe_new, mpe_destroy, NULL);
+            mpe_new, mpe_destroy, NULL);
     luaL_register(L, NULL, mpe_methods);
     lua_pop(L, 1);
 }
 
+static int mmi_new(lua_State *L) {
+    const char *arg = luaL_checkstring(L, 1);
+    const char *mdl = luaL_checkstring(L, 2);
+    const char *lat = luaL_checkstring(L, 3);
+    const char *ali = luaL_checkstring(L, 4);
+    KaldiMMI *mmi = new_KaldiMMI(arg, mdl, lat, ali);
+    luaT_pushudata(L, mmi, nerv_kaldi_mmi_tname);
+    return 1;
+}
+
+static int mmi_destroy(lua_State *L) {
+    KaldiMMI *mmi = luaT_checkudata(L, 1, nerv_kaldi_mmi_tname);
+    destroy_KaldiMMI(mmi);
+    return 0;
+}
+
+static int mmi_check(lua_State *L) {
+    KaldiMMI *mmi = luaT_checkudata(L, 1, nerv_kaldi_mmi_tname);
+    const Matrix *cumat = luaT_checkudata(L, 2, nerv_matrix_cuda_float_tname);
+    const char *utt = luaL_checkstring(L, 3);
+
+    lua_pushboolean(L, check_mmi(mmi, cumat, utt));
+    return 1;
+}
+
+static int mmi_calc_diff(lua_State *L) {
+    KaldiMMI *mmi = luaT_checkudata(L, 1, nerv_kaldi_mmi_tname);
+    Matrix *mat = luaT_checkudata(L, 2, nerv_matrix_host_float_tname);
+    const char *utt = luaL_checkstring(L, 3);
+
+    Matrix *diff = calc_diff_mmi(mmi, mat, utt);
+    luaT_pushudata(L, diff, nerv_matrix_host_float_tname);
+    return 1;
+}
+
+static int mmi_get_num_frames(lua_State *L) {
+    KaldiMMI *mmi = luaT_checkudata(L, 1, nerv_kaldi_mmi_tname);
+    lua_pushnumber(L, get_num_frames_mmi(mmi));
+    return 1;
+}
+
+static const luaL_Reg mmi_methods[] = {
+    {"check", mmi_check},
+    {"calc_diff", mmi_calc_diff},
+    {"get_num_frames", mmi_get_num_frames},
+    {NULL, NULL}
+};
+
+static void mmi_init(lua_State *L) {
+    luaT_newmetatable(L, nerv_kaldi_mmi_tname, NULL,
+            mmi_new, mmi_destroy, NULL);
+    luaL_register(L, NULL, mmi_methods);
+    lua_pop(L, 1);
+}
+
 void kaldi_seq_init(lua_State *L) {
     mpe_init(L);
+    mmi_init(L);
 }