class RustExecutor(BaseExecutor):
def __init__(self, model):
self.model_name = "score"
self.model = model
self.interpreter = RustInterpreter()
assembler_cls = get_assembler_cls(model)
self.model_ast = assembler_cls(model).assemble()
self.exec_path = None
def predict(self, X):
exec_args = [str(self.exec_path), *map(utils.format_arg, X)]
return utils.predict_from_commandline(exec_args)
def prepare(self):
if self.model_ast.output_size > 1:
execute_code = EXECUTE_VECTOR
else:
execute_code = EXECUTE_SCALAR
executor_code = EXECUTOR_CODE_TPL.format(
model_code=self.interpreter.interpret(self.model_ast),
execute_code=execute_code)
executor_file_name = self._resource_tmp_dir / f"{self.model_name}.rs"
utils.write_content_to_file(executor_code, executor_file_name)
self.exec_path = self._resource_tmp_dir / self.model_name
subprocess.call([
"rustc",
str(executor_file_name),
"-o",
str(self.exec_path)
])
def test_dependable_condition():
left = ast.BinNumExpr(
ast.IfExpr(
ast.CompExpr(ast.NumVal(1),
ast.NumVal(1),
ast.CompOpType.EQ),
ast.NumVal(1),
ast.NumVal(2)),
ast.NumVal(2),
ast.BinNumOpType.ADD)
right = ast.BinNumExpr(ast.NumVal(1), ast.NumVal(2), ast.BinNumOpType.DIV)
bool_test = ast.CompExpr(left, right, ast.CompOpType.GTE)
expr = ast.IfExpr(bool_test, ast.NumVal(1), ast.FeatureRef(0))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
let var0: f64;
let var1: f64;
if (1.0_f64) == (1.0_f64) {
var1 = 1.0_f64;
} else {
var1 = 2.0_f64;
}
if ((var1) + (2.0_f64)) >= ((1.0_f64) / (2.0_f64)) {
var0 = 1.0_f64;
} else {
var0 = input[0];
}
var0
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_log1p_expr():
expr = ast.Log1pExpr(ast.NumVal(2.0))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
f64::ln_1p(2.0_f64)
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_atan_expr():
expr = ast.AtanExpr(ast.NumVal(2.0))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
f64::atan(2.0_f64)
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_pow_expr():
expr = ast.PowExpr(ast.NumVal(2.0), ast.NumVal(3.0))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
f64::powf(2.0_f64, 3.0_f64)
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_raw_array():
expr = ast.VectorVal([ast.NumVal(3), ast.NumVal(4)])
expected_code = """
fn score(input: Vec<f64>) -> Vec<f64> {
vec![3.0_f64, 4.0_f64]
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_reused_expr():
reused_expr = ast.ExpExpr(ast.NumVal(1.0), to_reuse=True)
expr = ast.BinNumExpr(reused_expr, reused_expr, ast.BinNumOpType.DIV)
expected_code = """
fn score(input: Vec<f64>) -> f64 {
let var0: f64;
var0 = f64::exp(1.0_f64);
(var0) / (var0)
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_bin_num_expr():
expr = ast.BinNumExpr(
ast.BinNumExpr(
ast.FeatureRef(0), ast.NumVal(-2), ast.BinNumOpType.DIV),
ast.NumVal(2),
ast.BinNumOpType.MUL)
expected_code = """
fn score(input: Vec<f64>) -> f64 {
((input[0]) / (-2.0_f64)) * (2.0_f64)
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_sigmoid_expr():
expr = ast.SigmoidExpr(ast.NumVal(2.0))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
sigmoid(2.0_f64)
}
fn sigmoid(x: f64) -> f64 {
if x < 0.0_f64 {
let z: f64 = x.exp();
return z / (1.0_f64 + z);
}
1.0_f64 / (1.0_f64 + (-x).exp())
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_nested_condition():
left = ast.BinNumExpr(
ast.IfExpr(
ast.CompExpr(ast.NumVal(1),
ast.NumVal(1),
ast.CompOpType.EQ),
ast.NumVal(1),
ast.NumVal(2)),
ast.NumVal(2),
ast.BinNumOpType.ADD)
bool_test = ast.CompExpr(ast.NumVal(1), left, ast.CompOpType.EQ)
expr_nested = ast.IfExpr(bool_test, ast.FeatureRef(2), ast.NumVal(2))
expr = ast.IfExpr(bool_test, expr_nested, ast.NumVal(2))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
let var0: f64;
let var1: f64;
if (1.0_f64) == (1.0_f64) {
var1 = 1.0_f64;
} else {
var1 = 2.0_f64;
}
if (1.0_f64) == ((var1) + (2.0_f64)) {
let var2: f64;
if (1.0_f64) == (1.0_f64) {
var2 = 1.0_f64;
} else {
var2 = 2.0_f64;
}
if (1.0_f64) == ((var2) + (2.0_f64)) {
var0 = input[2];
} else {
var0 = 2.0_f64;
}
} else {
var0 = 2.0_f64;
}
var0
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_if_expr():
expr = ast.IfExpr(
ast.CompExpr(ast.NumVal(1), ast.FeatureRef(0), ast.CompOpType.EQ),
ast.NumVal(2),
ast.NumVal(3))
expected_code = """
fn score(input: Vec<f64>) -> f64 {
let var0: f64;
if (1.0_f64) == (input[0]) {
var0 = 2.0_f64;
} else {
var0 = 3.0_f64;
}
var0
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_bin_vector_num_expr():
expr = ast.BinVectorNumExpr(
ast.VectorVal([ast.NumVal(1), ast.NumVal(2)]),
ast.NumVal(1),
ast.BinNumOpType.MUL)
expected_code = """
fn score(input: Vec<f64>) -> Vec<f64> {
mul_vector_number(vec![1.0_f64, 2.0_f64], 1.0_f64)
}
fn add_vectors(v1: Vec<f64>, v2: Vec<f64>) -> Vec<f64> {
v1.iter().zip(v2.iter()).map(|(&x, &y)| x + y).collect::<Vec<f64>>()
}
fn mul_vector_number(v1: Vec<f64>, num: f64) -> Vec<f64> {
v1.iter().map(|&i| i * num).collect::<Vec<f64>>()
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_softmax_expr():
expr = ast.SoftmaxExpr([ast.NumVal(2.0), ast.NumVal(3.0)])
expected_code = """
fn score(input: Vec<f64>) -> Vec<f64> {
softmax(vec![2.0_f64, 3.0_f64])
}
fn softmax(x: Vec<f64>) -> Vec<f64> {
let size: usize = x.len();
let m: f64 = x.iter().fold(std::f64::MIN, |a, b| a.max(*b));
let mut exps: Vec<f64> = vec![0.0_f64; size];
let mut s: f64 = 0.0_f64;
for (i, &v) in x.iter().enumerate() {
exps[i] = (v - m).exp();
s += exps[i];
}
exps.iter().map(|&i| i / s).collect::<Vec<f64>>()
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)
def test_multi_output():
expr = ast.IfExpr(
ast.CompExpr(
ast.NumVal(1),
ast.NumVal(1),
ast.CompOpType.EQ),
ast.VectorVal([ast.NumVal(1), ast.NumVal(2)]),
ast.VectorVal([ast.NumVal(3), ast.NumVal(4)]))
expected_code = """
fn score(input: Vec<f64>) -> Vec<f64> {
let var0: Vec<f64>;
if (1.0_f64) == (1.0_f64) {
var0 = vec![1.0_f64, 2.0_f64];
} else {
var0 = vec![3.0_f64, 4.0_f64];
}
var0
}
"""
interpreter = RustInterpreter()
assert_code_equal(interpreter.interpret(expr), expected_code)