def get_tf_mod(verbose=False):
"""
:param bool verbose:
:return: module
"""
global _tf_mod
if _tf_mod:
return _tf_mod
import platform
from glob import glob
from returnn.tf.util.basic import OpCodeCompiler
# References:
# https://github.com/kpu/kenlm/blob/master/setup.py
# https://github.com/kpu/kenlm/blob/master/compile_query_only.sh
# Collect files.
assert kenlm_checked_out(), "submodule in %r not checked out?" % kenlm_dir
files = glob('%s/util/*.cc' % kenlm_dir)
files += glob('%s/lm/*.cc' % kenlm_dir)
files += glob('%s/util/double-conversion/*.cc' % kenlm_dir)
files = [fn for fn in files if not (fn.endswith('main.cc') or fn.endswith('test.cc'))]
assert files, "submodule in %r not checked out?" % kenlm_dir
libs = ["z"]
if platform.system() != 'Darwin':
libs.append('rt')
# Put code all together in one big blob.
src_code = ""
src_code += _kenlm_src_code_workarounds
for fn in files:
f_code = open(fn).read()
f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII
# We need to do some replacements to not clash symbol names.
fn_short = os.path.basename(fn).replace(".", "_")
for word in ["kConverter"]:
f_code = f_code.replace(word, "%s_%s" % (fn_short, word))
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(fn)
src_code += "\n\n// ------------ our code now: ------------\n\n"
src_code += "#line 1 \"<our code>\"\n"
src_code += _src_code
compiler = OpCodeCompiler(
base_name="KenLM", code_version=1, code=src_code,
include_paths=(kenlm_dir, kenlm_dir + "/util/double-conversion"),
c_macro_defines={"NDEBUG": 1, "KENLM_MAX_ORDER": 6, "HAVE_ZLIB": 1},
ld_flags=["-l%s" % lib for lib in libs],
is_cpp=True, use_cuda_if_available=False,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(tf_mod, "ken_lm_abs_score_strings"), "content of mod: %r" % (dir(tf_mod),)
_tf_mod = tf_mod
return tf_mod
def get_tf_mod(verbose=False):
"""
:param bool verbose:
:return: module
"""
global _tf_mod
if _tf_mod:
return _tf_mod
from glob import glob
from returnn.tf.util.basic import OpCodeCompiler
# Collect files.
assert openfst_checked_out(
), "submodule in %r not checked out?" % openfst_dir
files = glob('%s/src/lib/*.cc' % openfst_dir)
assert files, "submodule in %r not checked out?" % openfst_dir
files = sorted(files) # make somewhat deterministic
libs = []
if platform.system() != 'Darwin':
libs.append('rt')
# Put code all together in one big blob.
src_code = ""
for fn in files:
f_code = open(fn).read()
f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(
fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(
fn)
src_code += "\n\n// ------------ our code now: ------------\n\n"
src_code += "#line 1 \"returnn/tf/util/open_fst.py:_src_code\"\n"
src_code += _src_code
compiler = OpCodeCompiler(
base_name="OpenFst",
code_version=1,
code=src_code,
include_paths=("%s/src/include" % openfst_dir, ),
c_macro_defines={
"NDEBUG":
1, # https://github.com/tensorflow/tensorflow/issues/17316
},
ld_flags=["-l%s" % lib for lib in libs],
is_cpp=True,
use_cuda_if_available=False,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(
tf_mod, "open_fst_transition"), "content of mod: %r" % (dir(tf_mod), )
_tf_mod = tf_mod
return tf_mod
def init_warprnnt(verbose=False):
"""
Initializes and compiles the library. Caches the TF module.
:param bool verbose:
"""
global _tf_mod
if _tf_mod:
return
assert is_checked_out(
), "submodule not checked out? Run `git submodule update --init --recursive`"
# References:
# https://github.com/HawkAaron/warp-transducer/blob/master/tensorflow_binding/setup.py
src_files = [
'%s/src/rnnt_entrypoint.cpp' % submodule_dir,
'%s/tensorflow_binding/src/warprnnt_op.cc' % submodule_dir
]
assert all([os.path.isfile(f) for f in src_files
]), "submodule in %r not checked out?" % warprnnt_dir
src_code = ""
for fn in src_files:
f_code = open(fn).read()
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(
fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(
fn)
with_cuda = CudaEnv.get_instance().is_available() and is_gpu_available()
with_omp = sys.platform.startswith("linux")
compiler = OpCodeCompiler(
base_name="warprnnt_kernels",
code_version=1,
code=src_code,
include_paths=(submodule_dir + "/include", ),
c_macro_defines=dict_joined(
{"WITH_OMP": 1} if with_omp else {"RNNT_DISABLE_OMP": 1},
{"WARPRNNT_ENABLE_GPU": 1} if with_cuda else {}),
ld_flags=["-Xcompiler", "-fopenmp"] if with_omp else [],
is_cpp=True,
use_cuda_if_available=True,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(tf_mod, "WarpRNNT"), "content of mod: %r" % (dir(tf_mod), )
_tf_mod = tf_mod
return tf_mod
def init_warprna(verbose=False):
"""
Initializes and compiles the library. Caches the TF module.
:param bool verbose:
"""
global _tf_mod
if _tf_mod:
return
assert is_checked_out(
), "submodule not checked out? Run `git submodule update --init --recursive`"
assert is_tf_cuda_build(), "TF with CUDA support required"
src_files = [
'%s/tensorflow_binding/src/warp_rna_op.cc' % submodule_dir,
'%s/core.cu' % submodule_dir,
]
src_code = ""
for fn in src_files:
f_code = open(fn).read()
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(
fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(
fn)
compiler = OpCodeCompiler(base_name="warprna_kernels",
code_version=1,
code=src_code,
include_paths=(submodule_dir,
os.path.dirname(submodule_dir)),
c_macro_defines={"WARPRNA_ENABLE_GPU": 1},
is_cpp=True,
use_cuda_if_available=True,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(tf_mod, "WarpRNA"), "content of mod: %r" % (dir(tf_mod), )
_tf_mod = tf_mod
return tf_mod
def get_filtered_score_op(verbose=False):
cpp_code = """
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/resource_op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/public/version.h"
#include <cmath>
#include <map>
#include <set>
#include <string>
#include <tuple>
using namespace tensorflow;
REGISTER_OP("GetFilteredScore")
.Input("prev_str: string")
.Input("scores: float32")
//.Input("labels: string")
.Output("new_scores: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(1));
return Status::OK();
});
class GetFilteredScoreOp : public OpKernel {
public:
using OpKernel::OpKernel;
void Compute(OpKernelContext* context) override {
const Tensor* prev_str = &context->input(0);
const Tensor* scores = &context->input(1);
//const Tensor* labels = &context->input(2);
int n_batch = prev_str->shape().dim_size(0);
int n_beam = prev_str->shape().dim_size(1);
Tensor* ret;
OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape({n_batch, n_beam}), &ret));
for(int bat = 0; bat < n_batch; ++bat)
for(int hyp = 0; hyp < n_beam; ++hyp)
ret->tensor<float, 2>()(bat, hyp) = scores->tensor<float, 2>()(bat, hyp);
for(int bat = 0; bat < n_batch; ++bat) {
std::map<tstring, std::set<int> > new_hyps; // seq -> set of hyp idx
for(int hyp = 0; hyp < n_beam; ++hyp) {
auto& seq_set = new_hyps[prev_str->tensor<tstring, 2>()(bat, hyp)];
seq_set.insert(hyp);
}
for(const auto& items : new_hyps) {
if(std::get<1>(items).size() > 1) {
float best_score = 0.;
int best_idx = -1;
for(int idx : std::get<1>(items)) {
float score = scores->tensor<float, 2>()(bat, idx);
if(score > best_score || best_idx == -1) {
best_score = score;
best_idx = idx;
}
}
float sum_score = 0.;
for(int idx : std::get<1>(items)) {
float score = scores->tensor<float, 2>()(bat, idx);
sum_score += expf(score - best_score);
}
sum_score = logf(sum_score) + best_score;
for(int idx : std::get<1>(items)) {
if(idx != best_idx)
ret->tensor<float, 2>()(bat, idx) = -std::numeric_limits<float>::infinity();
else
ret->tensor<float, 2>()(bat, idx) = sum_score;
}
}
}
}
}
};
REGISTER_KERNEL_BUILDER(Name("GetFilteredScore").Device(DEVICE_CPU), GetFilteredScoreOp);
"""
from returnn.tf.util.basic import OpCodeCompiler
compiler = OpCodeCompiler(base_name="GetFilteredScore",
code_version=1,
code=cpp_code,
is_cpp=True,
use_cuda_if_available=False,
verbose=verbose)
tf_mod = compiler.load_tf_module()
return tf_mod.get_filtered_score