Skip to content
Snippets Groups Projects
linkprediction.py 1.7 KiB
Newer Older
#!/usr/bin/env python
import os
import sys
import inspect
import networkx as nx
import json

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir) 

from MLBasicFunctions import addNewEdgeMetaData, buildGraph
from MLRepositoryInterface import MLServerInterface

CONST_MLQUEUENAME = "lpr_queue"

class MLServer(MLServerInterface):
    def __init__(self):
        # Fill in the parameters for communication with the algorithm here
        self.parameters = {}

        self.this = {}
        self.this[CONST_MLQUEUENAME] = self.parameters

    def decodeMessage(self, data):
        graph = buildGraph(data)
        return graph

    # Adjust the data to fit our existing format of edges for easier future visualisation
    def adjustData(self, mlresult):
        list = []    
        for index,(fr, to, value) in enumerate(mlresult):
            dict = {}
            dict["attributes"] = {}
            dict["attributes"]["jaccard_coefficient"] = value
            dict["from"] = fr
            dict["id"] = "link_prediction_relation/" + str(index)
            dict["to"] = to  
            list.append(dict)   
        return list

    def handleJSON(self, body):
        data = json.loads(body.decode())
        print(data["machineLearning"][0]["params"])
        print(data["queryID"])
        print(data["type"])
        G = self.decodeMessage(data)
        mlresult = nx.jaccard_coefficient(G)
        mlresultparsed = self.adjustData(mlresult)
        result = addNewEdgeMetaData(data,mlresultparsed)
        #print(result)

        #Transforms the result (a dictionary) into JSON
        return json.dumps(result, indent = 4)