import numpy as np
from scipy.spatial.distance import cosine
import logging
from multiprocessing import Pool


def xcov(set1, set2):
    num_set1, M = set1.shape
    num_set2, N = set2.shape
    # one-to-one paired
    assert(num_set1 == num_set2)

    cov = np.dot(np.transpose(set1), set2)
    logging.info("get XCOV matrix of shape R^%d x %d" % (cov.shape[0], cov.shape[1]))
    return cov


def train(Qs, As):
    '''
    params q: sentence embedding for question set
    params a: sentence embedding for answer set
    '''
    if isinstance(Qs, list):
        Qs = np.asarray(Qs, dtype="float64")
    if isinstance(As, list):
        As = np.asarray(As, dtype="float64")

    logging.info("computing cross-covariance matrix")
    cov = xcov(Qs, As)
    logging.info("decomposition using SVD")
    U, s, V = np.linalg.svd(cov, full_matrices=False)
    return U, V


def find_answer(v_q, As):
    pool = Pool(processes=8)

    # get distance between the question and answer, return with the answer index
    def distance(indx_t):
        indx, t = indx_t
        dist = cosine(v_q, t)
        return indx, dist

    result = pool.map(distance, As)
    best_indx, _ = min(result, key=lambda x: x[1])

    return best_indx

@@ -102,16 +102,46 @@
 rt Pool%0A
+from functools import partial%0A
 %0A%0Adef xc
@@ -897,68 +897,8 @@
 V%0A%0A%0A
-def find_answer(v_q, As):%0A    pool = Pool(processes=8)%0A%0A    
 # ge
@@ -970,20 +970,16 @@
 r index%0A
-    
 def dist
@@ -983,16 +983,21 @@
 istance(
+v_q, 
 indx_t):
@@ -993,28 +993,24 @@
 q, indx_t):%0A
-    
     indx, t 
@@ -1018,20 +1018,16 @@
  indx_t%0A
-    
     dist
@@ -1048,20 +1048,16 @@
  t)%0A    
- 
-   
 return i
@@ -1067,16 +1067,83 @@
 , dist%0A%0A
+%0Adef find_answer(v_q, As):%0A    with Pool(processes=8) as pool:%0A    
     resu
@@ -1156,16 +1156,24 @@
 ool.map(
+partial(
 distance
@@ -1173,16 +1173,26 @@
 istance,
+ v_q=v_q),
  As)%0A   

fix bug