code for creation of triplet loss sampels

create_tripet_loss_samples.py

import itertools

import numpy as np

from diskdict import DiskDict
from diskarray import DiskArray
from basescript import BaseScript
from wordvecspace import WordVecSpaceMem

class TripletLossSamples(BaseScript):
    def __init__(self):
        super(TripletLossSamples, self).__init__()
        self.train_d = DiskArray(self.args.out_f, shape=(0,), dtype=self._get_dtype())
        self.wv = WordVecSpaceMem(self.args.wvspace)
        self.d = DiskDict(self.args.clusters_f)

    def _get_dtype(self):
        d_type = [
                    ('anchor', np.float32, 300),
                    ('positive', np.float32, 300),
                    ('negative', np.float32, 300),
                ]

        return d_type

    def create_triples(self, positives, negatives, anchor):
        for i in range(len(positives)):
            p = '<id:' + positives[i] + '>'
            for i in range(len(negatives)):
                n = '<id:' + negatives[i] + '>'
                a_v = self.wv.get_word_vector(anchor)
                p_v = self.wv.get_word_vector(p)
                n_v = self.wv.get_word_vector(n)
                self.train_d.append((a_v, p_v, n_v))
  
    def run(self):
        for k, cluster in self.d.items():  
            positives = cluster['positive']
            negatives = cluster['nagative'] #FIXME Typo error in diskdict file.
            anchor = cluster['anchor']
             self.create_triples(positives, negatives, anchor)
 
             for i in range(len(positives)):
                 _anchor = '<id:' + positives[i] + '>'
                 if anchor == _anchor:
                     continue
 
                 self.create_triples(positives, negatives, _anchor)
 
        self.train_d.flush()
 
     def define_args(self, parser):
         parser.add_argument('wvspace', help='vector space file')
         parser.add_argument('clusters_f', help='diskdict file which contains ids')
         parser.add_argument('out_f', help='output file to store labels')
 
 if __name__ == '__main__':
     TripletLossSamples().start()