from models.gae import GAE
from models.resource_embedding import ResourceEmbedding
from utils.data_processing import single_graph_to_pyg_data


class ResourceEmbeddingDvfs(ResourceEmbedding):

    def run(self, graph):
        g = single_graph_to_pyg_data(graph,
                                     lambda x: x['features'],
                                     self.num_features)
        return super().get_embeddings(g.x, g.edge_index, g.edge_weight)