Nearest Neighbour retriever

tf_retriever.py

class TFRanker:
    def __init__(self, dim, metric, top_k=3):
        self.dim = dim
        self.top_k = top_k
        self.metric = metric
        self.graph = tf.Graph()
        self.session = tf.Session(graph=self.graph)
        self.build_graph()
        
        
    def build_graph(self):
        with self.graph.as_default():
            self.query = tf.placeholder("float", [self.dim])
            self.kbase = tf.placeholder("float", [None, self.dim])

            distance = self.metric(self.kbase, self.query)

            top_neg_dists, top_indices = tf.math.top_k(tf.negative(distance), k=self.top_k)
            top_dists = tf.negative(top_neg_dists)

            self.top_distances = top_dists
            self.top_indices = top_indices
        
    def predict(self, kbase, query):
        with self.graph.as_default():
            I, D = self.session.run([self.top_indices, self.top_distances],
                                    feed_dict={self.query: query, self.kbase: kbase})
        return I, D
    
def euclidean_distance(kbase, query):
    sqr_distance = tf.reduce_sum(
        tf.pow(kbase - query, 2),
        reduction_indices=1)
    distance = tf.sqrt(sqr_distance)
    return distance