def get_big_model_nodes(orig_model_node: WillumpModelNode, new_input_node: WillumpGraphNode,
small_model_output_node: CascadeThresholdProbaNode, small_model_output_name: str) \
-> Tuple[WillumpModelNode, CascadeCombinePredictionsNode]:
assert (isinstance(orig_model_node, WillumpPredictNode)
or (top_k is not None
and isinstance(orig_model_node, WillumpPredictProbaNode)))
assert (len(new_input_node.get_output_names()) == 1)
output_name = orig_model_node.get_output_name()
output_type = orig_model_node.get_output_type()
new_input_name = new_input_node.get_output_names()[0]
if top_k is None:
big_model_output = WillumpPredictNode(
model_name=BIG_MODEL_NAME,
x_name=new_input_name,
x_node=new_input_node,
output_name=output_name,
output_type=output_type,
input_width=orig_model_node.input_width)
else:
big_model_output = WillumpPredictProbaNode(
model_name=BIG_MODEL_NAME,
x_name=new_input_name,
x_node=new_input_node,
output_name=output_name,
output_type=output_type,
input_width=orig_model_node.input_width)
combining_node = CascadeCombinePredictionsNode(
big_model_predictions_node=big_model_output,
big_model_predictions_name=output_name,
small_model_predictions_node=small_model_output_node,
small_model_predictions_name=small_model_output_name,
output_name=output_name,
output_type=output_type)
return big_model_output, combining_node
def get_small_model_nodes(orig_model_node: WillumpModelNode, new_input_node: WillumpGraphNode) \
-> Tuple[WillumpModelNode, WillumpGraphNode]:
assert (isinstance(orig_model_node, WillumpPredictNode)
or (top_k is not None
and isinstance(orig_model_node, WillumpPredictProbaNode)))
assert (len(new_input_node.get_output_names()) == 1)
proba_output_name = "small__proba_" + orig_model_node.get_output_name()
output_type = WeldVec(WeldDouble())
typing_map[proba_output_name] = output_type
new_input_name = new_input_node.get_output_names()[0]
predict_proba_node = WillumpPredictProbaNode(
model_name=SMALL_MODEL_NAME,
x_name=new_input_name,
x_node=new_input_node,
output_name=proba_output_name,
output_type=output_type,
input_width=orig_model_node.input_width)
threshold_output_name = "small_preds_" + orig_model_node.get_output_name(
)
if top_k is None:
threshold_node = CascadeThresholdProbaNode(
input_node=predict_proba_node,
input_name=proba_output_name,
output_name=threshold_output_name,
threshold=cascade_threshold)
else:
threshold_node = CascadeTopKSelectionNode(
input_node=predict_proba_node,
input_name=proba_output_name,
output_name=threshold_output_name,
top_k=top_k * cascade_threshold)
typing_map[threshold_output_name] = WeldVec(WeldChar())
return predict_proba_node, threshold_node
def analyze_return(self, node: ast.Return) -> None:
output_name: str = "__willump_return_result"
function_name: str = node.value.func.id
assert (len(node.value.args) == 2) # First is the model parameter, second is a list of model arguments.
model_param: str = node.value.args[0].id
input_names, input_nodes = [], []
for function_argument in node.value.args[1].elts:
input_name = function_argument.id
input_names.append(input_name)
if input_name in self._variable_dict:
input_nodes.append(self._variable_dict[input_name])
else:
input_nodes.append(None)
model_node: WillumpGraphNode = WillumpGraphNode(function_name=function_name,
output_name=output_name,
input_names=input_names,
input_nodes=input_nodes,
model_param=model_param,
cost=0)
self._model_node = model_node
def analyze_assign(self, node: ast.Assign) -> None:
assert (len(node.targets) == 1) # Assume assignment to only one variable.
output_name: str = node.targets[0].id
function_name: str = node.value.func.id
input_names, input_nodes = [], []
for function_argument in node.value.args:
input_name = function_argument.id
input_names.append(input_name)
if input_name in self._variable_dict:
input_nodes.append(self._variable_dict[input_name])
else:
input_nodes.append(None)
keywords = node.value.keywords
function_node: WillumpGraphNode = WillumpGraphNode(function_name=function_name,
output_name=output_name,
input_names=input_names,
input_nodes=input_nodes,
keywords=keywords,
cost=self._timing_map[output_name])
self._variable_dict[output_name] = function_node
def push_cascade(self, small_model_output_node: WillumpGraphNode):
small_model_output_name = small_model_output_node.get_output_names()[0]
self._small_model_output_name = small_model_output_name
self._input_nodes.append(small_model_output_node)
self._input_names.append(small_model_output_name)
def get_combiner_node_eval(mi_head: WillumpGraphNode, li_head: WillumpGraphNode, orig_node: WillumpGraphNode,
small_model_output_node: CascadeThresholdProbaNode) \
-> WillumpGraphNode:
"""
Generate a node that will fuse node_one and node_two to match the output of orig_node. Requires all
three nodes be of the same type.
"""
assert (len(orig_node.get_output_types()) == len(
mi_head.get_output_types()) == len(li_head.get_output_types()) ==
1)
orig_output_type = orig_node.get_output_types()[0]
if isinstance(orig_output_type, WeldCSR):
return CascadeStackSparseNode(
more_important_nodes=[mi_head],
more_important_names=[mi_head.get_output_names()[0]],
less_important_nodes=[li_head],
less_important_names=[li_head.get_output_names()[0]],
small_model_output_node=small_model_output_node,
small_model_output_name=small_model_output_node.
get_output_name(),
output_name=orig_node.get_output_names()[0],
output_type=orig_output_type)
elif isinstance(orig_output_type, WeldPandas):
mi_output_types = mi_head.get_output_types()[0]
li_output_types = li_head.get_output_types()[0]
assert (isinstance(mi_output_types, WeldPandas)
and isinstance(li_output_types, WeldPandas))
new_selected_columns = mi_output_types.column_names + li_output_types.column_names
new_field_types = mi_output_types.field_types + li_output_types.field_types
output_name = orig_node.get_output_names()[0]
new_output_type = WeldPandas(column_names=new_selected_columns,
field_types=new_field_types)
typing_map[output_name] = new_output_type
return CascadeColumnSelectionNode(
more_important_nodes=[mi_head],
more_important_names=[mi_head.get_output_names()[0]],
more_important_types=[mi_output_types],
less_important_nodes=[li_head],
less_important_names=[li_head.get_output_names()[0]],
less_important_types=[li_output_types],
output_name=output_name,
small_model_output_node=small_model_output_node,
small_model_output_name=small_model_output_node.
get_output_name(),
selected_columns=new_selected_columns)
elif isinstance(orig_output_type, WeldVec):
assert (isinstance(orig_output_type.elemType, WeldVec))
return CascadeStackDenseNode(
more_important_nodes=[mi_head],
more_important_names=[mi_head.get_output_names()[0]],
less_important_nodes=[li_head],
less_important_names=[li_head.get_output_names()[0]],
small_model_output_node=small_model_output_node,
small_model_output_name=small_model_output_node.
get_output_name(),
output_name=orig_node.get_output_names()[0],
output_type=orig_output_type)
else:
panic("Unrecognized eval combiner output type %s" %
orig_output_type)