def graph_encoder(
populated_graph_db: graph_tuple_database.Database,
ir2seq_encoder: ir2seq.EncoderBase,
cache_size: int,
):
"""A test fixture which enumerates statement encoders."""
return graph2seq.GraphEncoder(populated_graph_db, ir2seq_encoder,
cache_size)
def GetEncoder(self) -> graph2seq.EncoderBase:
"""Construct the graph encoder."""
if not (self.graph_db.graph_y_dimensionality
and self.graph_db.node_x_dimensionality == 2
and self.graph_db.node_y_dimensionality == 0):
raise app.UsageError(
f"Unsupported graph dimensionalities: {self.graph_db}")
return graph2seq.GraphEncoder(graph_db=self.graph_db,
ir2seq_encoder=self._ir2seq_encoder)