exputils.py

import numpy as np


def get_unique_triplets(triplets, y):
    unique_triplets = set()

    unique_triplets_list = []
    unique_y_list = []

    for triplet_i, y_i in zip(triplets, y):
        triplet_i_tuple = tuple(triplet_i)
        if triplet_i_tuple in unique_triplets:
            continue
        unique_triplets.add(triplet_i_tuple)
        unique_triplets_list.append(triplet_i)
        unique_y_list.append(y_i)
    return np.array(unique_triplets_list), np.array(unique_y_list)


def get_unique_items(triplets):
    unique_items = set()

    for triplet in triplets:
        for item in triplet:
            unique_items.add(str(item))

    return sorted(list(unique_items))


def get_items_in_triplets(triplets):
    items = set()
    for triplet in triplets:
        for item in triplet:
            items.add(item.strip())
    return items


def get_triplets_emb(triplets, items, phi_emb):
    emb_map = {}
    for k, v in zip(items, phi_emb):
        emb_map[str(k)] = v

    triplets_emb = []
    for row in triplets:
        triplets_emb.append(
            [
                emb_map[str(row[0]).strip()],
                emb_map[str(row[1]).strip()],
                emb_map[str(row[2]).strip()],
            ]
        )
    return np.array(triplets_emb)


def get_triplets_with_items(triplets, items, y):
    items = set(items)
    triplets_with_items = []
    y_with_items = []
    triplets_without_items = []
    y_without_items = []

    for triplet, y_i in zip(triplets, y):
        if triplet[0] in items or triplet[1] in items or triplet[2] in items:
            triplets_with_items.append(triplet)
            y_with_items.append(y_i)
        else:
            triplets_without_items.append(triplet)
            y_without_items.append(y_i)

    return triplets_with_items, triplets_without_items, y_with_items, y_without_items