gpt2_embed_relatedness.py

def sim(ta, tb):
    a_t = enc.encode(ta)
    b_t = enc.encode(tb)
    a_te = np.array(model.transformer.wte.weight[a_t].tolist())
    b_te = np.array(model.transformer.wte.weight[b_t].tolist())
    a_pe = np.array(model.transformer.wpe.weight[[i for i in range(len(a_t))]].tolist())
    b_pe = np.array(model.transformer.wpe.weight[[i for i in range(len(b_t))]].tolist())
    x = np.add.reduce(a_te+a_pe)
    y = np.add.reduce(b_te+b_pe)
    
    similarity = 1 - spatial.distance.cosine(x, y)
    
    return similarity