load_data.py

__author__ = 'thiagocastroferreira'

"""
Author: Thiago Castro Ferreira
Date: 12/12/2017
Description:
    Script for loading the referring expressions collection.

    PYTHON VERSION: 2.7

    UPDATE CONSTANTS:
        VOCAB_PATH
        TRAIN_REFEX_PATH
        DEV_REFEX_PATH
        TEST_REFEX_PATH
"""

import os

# PATH FOR VOCABULARY
VOCAB_PATH = 'data/'

# PATH FOR REFERRING EXPRESSION COLLECTIONS
TRAIN_REFEX_PATH = 'data/train'
DEV_REFEX_PATH = 'data/dev'
TEST_REFEX_PATH = 'data/test'

def load(fpre_context, fpos_context, fentity, frefex, fsize, character):
    with open(fpre_context) as f:
        pre_context = map(lambda x: x.split(), f.read().split('\n'))

    with open(fpos_context) as f:
        pos_context = map(lambda x: x.split(), f.read().split('\n'))

    with open(fentity) as f:
        entity = f.read().split('\n')

    with open(frefex) as f:
        if character:
            refex = map(lambda x: ['eos'] + list(x.replace('eos', '').strip()) + ['eos'], f.read().split('\n'))
        else:
            refex = map(lambda x: x.split(), f.read().split('\n'))

    with open(fsize) as f:
        size = f.read().split('\n')

    return {
        'pre_context':list(pre_context),
        'pos_context':list(pos_context),
        'entity':list(entity),
        'refex':list(refex),
        'size':list(size)
    }

def run(character=False):
    # VOCABULARY
    with open(os.path.join(VOCAB_PATH, 'input_vocab.txt')) as f:
        input_vocab = f.read().split('\n')

    if character:
        with open(os.path.join(VOCAB_PATH, 'character_vocab.txt')) as f:
            output_vocab = f.read().split('\n')
    else:
        with open(os.path.join(VOCAB_PATH, 'output_vocab.txt')) as f:
            output_vocab = f.read().split('\n')
    vocab = {'input':input_vocab, 'output':output_vocab}

    # TRAINSET
    fprecontext = os.path.join(TRAIN_REFEX_PATH, 'pre_context.txt')
    fposcontext = os.path.join(TRAIN_REFEX_PATH, 'pos_context.txt')
    fentity = os.path.join(TRAIN_REFEX_PATH, 'entity.txt')
    frefex = os.path.join(TRAIN_REFEX_PATH, 'refex.txt')
    fsize = os.path.join(TRAIN_REFEX_PATH, 'size.txt')
    trainset = load(fprecontext, fposcontext, fentity, frefex, fsize, character)

    # DEVSET
    fprecontext = os.path.join(DEV_REFEX_PATH, 'pre_context.txt')
    fposcontext = os.path.join(DEV_REFEX_PATH, 'pos_context.txt')
    fentity = os.path.join(DEV_REFEX_PATH, 'entity.txt')
    frefex = os.path.join(DEV_REFEX_PATH, 'refex.txt')
    fsize = os.path.join(DEV_REFEX_PATH, 'size.txt')
    devset = load(fprecontext, fposcontext, fentity, frefex, fsize, character)

    # TESTSET
    fprecontext = os.path.join(TEST_REFEX_PATH, 'pre_context.txt')
    fposcontext = os.path.join(TEST_REFEX_PATH, 'pos_context.txt')
    fentity = os.path.join(TEST_REFEX_PATH, 'entity.txt')
    frefex = os.path.join(TEST_REFEX_PATH, 'refex.txt')
    fsize = os.path.join(TEST_REFEX_PATH, 'size.txt')
    testset = load(fprecontext, fposcontext, fentity, frefex, fsize, character)

    return vocab, trainset, devset, testset