<gh_stars>10-100
import argparse
import numpy as np
import pickle
import random
import os
from ogb.lsc import MAG240MDataset
from collections import defaultdict
from cleora import train_cleora
from tqdm import tqdm
from coders import get_vcoder
from root import ROOT


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--working-dir", type=str, default = 'data', help='Working directory')
    parser.add_argument("--sketch-dim", type=int, default = 128, help='Sketch width')
    parser.add_argument("--sketch-depth", type=int, default = 40, help='Sketch depth')
    parser.add_argument("--dim", type=int, default = 512, help='Node embedding dimension')
    return parser


def create_cleora_input_file(filename, paper2inst):
    print("Creating Cleora input file")
    with open(filename, 'w') as f:
        for paper, inst in paper2inst.items():
            if len(inst) > 1:
                if len(inst) > 20:
                    inst = random.sample(inst, 20)
                f.write(' '.join([str(i) for i in inst]))
                f.write('\n')


def get_cleora_embeddings(filename):
    """
    Read embeddings from file generated by Cleora.
    """
    embeddings = []
    ids = []
    with open(filename, 'r') as f:
        next(f)
        for index, line in enumerate(f):
            line_splitted = line.split(sep=' ')
            ids.append(line_splitted[0])
            embeddings.append(list(map(float, line_splitted[2:])))
    embeddings = np.array(embeddings)
    return ids, embeddings


def compute_sketches(params):
    os.makedirs(params.working_dir, exist_ok=True)
    edge_author_paper = np.load(f'{ROOT}/mag240m_kddcup2021/processed/author___writes___paper/edge_index.npy', mmap_mode='r')
    edges_author_inst = np.load(f'{ROOT}/mag240m_kddcup2021/processed/author___affiliated_with___institution/edge_index.npy',  mmap_mode='r')
    print("Creating mapping author -> institutions")
    author2inst = defaultdict(set)
    for i in tqdm(range(edges_author_inst.shape[1])):
        author = edges_author_inst[0,i]
        inst = edges_author_inst[1,i]
        author2inst[author].add(inst)

    dataset = MAG240MDataset(ROOT)
    split_dict = dataset.get_idx_split()
    train_idx = split_dict['train']
    valid_idx = split_dict['valid']
    test_idx = split_dict['test']

    idxs_labelled = set(train_idx)
    idxs_labelled = idxs_labelled.union(set(valid_idx))
    idxs_labelled = idxs_labelled.union(set(test_idx))

    print("Creating mapping paper -> institutions")
    paper2inst = defaultdict(set)

    for i in tqdm(range(edge_author_paper.shape[1])):
        author =  edge_author_paper[0,i]
        paper = edge_author_paper[1,i]
        paper_inst = author2inst[author]
        for inst in paper_inst:
            paper2inst[paper].add(inst)


    input_filename = f"{params.working_dir}/cleora_input"
    create_cleora_input_file(input_filename, paper2inst)
    train_cleora(dim=params.dim, iter_=3, columns='complex::reflexive::institution', input_filename=input_filename,
             working_dir=params.working_dir)

    output_filename = f'{params.working_dir}/emb__institution__institution.out'
    ids, embeddings = get_cleora_embeddings(output_filename)
    vcoder = get_vcoder(embeddings, params.sketch_depth, params.sketch_dim)

    codes = vcoder.transform(embeddings).astype(np.uint8)
    np.save(f'{params.working_dir}/inst_codes.npy', codes)

    inst_codes_inst2id = {inst:i for i, inst in enumerate(ids)}
    with open(f'{params.working_dir}/codes_inst2id', 'wb') as handle:
        pickle.dump(inst_codes_inst2id, handle, protocol=pickle.HIGHEST_PROTOCOL)

    paper2inst_label = {k: paper2inst[k] for k in idxs_labelled}
    with open(f'{params.working_dir}/paper2inst', 'wb') as handle:
        pickle.dump(paper2inst_label, handle, protocol=pickle.HIGHEST_PROTOCOL)

if __name__ == "__main__":
    parser = get_parser()
    params = parser.parse_args()
    compute_sketches(params)
import argparse
import numpy as np
import pickle
import random
import os
from ogb.lsc import MAG240MDataset
from collections import defaultdict
from cleora import train_cleora
from tqdm import tqdm
from coders import get_vcoder
from root import ROOT


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--working-dir", type=str, default = 'data', help='Working directory')
    parser.add_argument("--sketch-dim", type=int, default = 128, help='Sketch width')
    parser.add_argument("--sketch-depth", type=int, default = 40, help='Sketch depth')
    parser.add_argument("--dim", type=int, default = 512, help='Node embedding dimension')
    return parser


def create_cleora_input_file(filename, paper2inst):
    print("Creating Cleora input file")
    with open(filename, 'w') as f:
        for paper, inst in paper2inst.items():
            if len(inst) > 1:
                if len(inst) > 20:
                    inst = random.sample(inst, 20)
                f.write(' '.join([str(i) for i in inst]))
                f.write('\n')


def get_cleora_embeddings(filename):
    """
    Read embeddings from file generated by Cleora.
    """
    embeddings = []
    ids = []
    with open(filename, 'r') as f:
        next(f)
        for index, line in enumerate(f):
            line_splitted = line.split(sep=' ')
            ids.append(line_splitted[0])
            embeddings.append(list(map(float, line_splitted[2:])))
    embeddings = np.array(embeddings)
    return ids, embeddings


def compute_sketches(params):
    os.makedirs(params.working_dir, exist_ok=True)
    edge_author_paper = np.load(f'{ROOT}/mag240m_kddcup2021/processed/author___writes___paper/edge_index.npy', mmap_mode='r')
    edges_author_inst = np.load(f'{ROOT}/mag240m_kddcup2021/processed/author___affiliated_with___institution/edge_index.npy',  mmap_mode='r')
    print("Creating mapping author -> institutions")
    author2inst = defaultdict(set)
    for i in tqdm(range(edges_author_inst.shape[1])):
        author = edges_author_inst[0,i]
        inst = edges_author_inst[1,i]
        author2inst[author].add(inst)

    dataset = MAG240MDataset(ROOT)
    split_dict = dataset.get_idx_split()
    train_idx = split_dict['train']
    valid_idx = split_dict['valid']
    test_idx = split_dict['test']

    idxs_labelled = set(train_idx)
    idxs_labelled = idxs_labelled.union(set(valid_idx))
    idxs_labelled = idxs_labelled.union(set(test_idx))

    print("Creating mapping paper -> institutions")
    paper2inst = defaultdict(set)

    for i in tqdm(range(edge_author_paper.shape[1])):
        author =  edge_author_paper[0,i]
        paper = edge_author_paper[1,i]
        paper_inst = author2inst[author]
        for inst in paper_inst:
            paper2inst[paper].add(inst)


    input_filename = f"{params.working_dir}/cleora_input"
    create_cleora_input_file(input_filename, paper2inst)
    train_cleora(dim=params.dim, iter_=3, columns='complex::reflexive::institution', input_filename=input_filename,
             working_dir=params.working_dir)

    output_filename = f'{params.working_dir}/emb__institution__institution.out'
    ids, embeddings = get_cleora_embeddings(output_filename)
    vcoder = get_vcoder(embeddings, params.sketch_depth, params.sketch_dim)

    codes = vcoder.transform(embeddings).astype(np.uint8)
    np.save(f'{params.working_dir}/inst_codes.npy', codes)

    inst_codes_inst2id = {inst:i for i, inst in enumerate(ids)}
    with open(f'{params.working_dir}/codes_inst2id', 'wb') as handle:
        pickle.dump(inst_codes_inst2id, handle, protocol=pickle.HIGHEST_PROTOCOL)

    paper2inst_label = {k: paper2inst[k] for k in idxs_labelled}
    with open(f'{params.working_dir}/paper2inst', 'wb') as handle:
        pickle.dump(paper2inst_label, handle, protocol=pickle.HIGHEST_PROTOCOL)

if __name__ == "__main__":
    parser = get_parser()
    params = parser.parse_args()
    compute_sketches(params)
Read embeddings from file generated by Cleora.