from ..dataset import Cog_Dataset


class WIKIDATA5MProcessor:
    def __init__(self, lut_E, lut_R):
        self.lut_E = lut_E
        self.lut_R = lut_R

    def process(self, datable):
        datable = self.str2number(datable)
        dataset = Cog_Dataset(data=datable, task="kr")
        return dataset

    def str2number(self, datable):
        for i in range(len(datable)):
            datable["head"][i] = self.lut_E.str_dic[datable["head"][i]]
            datable["relation"][i] = self.lut_R.str_dic[datable["relation"][i]]
            datable["tail"][i] = self.lut_E.str_dic[datable["tail"][i]]
        return datable
