# coding=utf-8

# Copyright 2017 The Tensor2Tensor Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Data generators for the Description2Code OpenAI data-set."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import collections

import os

import random

import re

import zipfile

# Dependency imports

from tensor2tensor.data_generators import generator_utils

from tensor2tensor.data_generators import problem

from tensor2tensor.data_generators import text_encoder

from tensor2tensor.utils import registry

import tensorflow as tf

# End-of-sentence marker.

EOS = text_encoder.EOS_ID

_DATASET_URL = "https://drive.google.com/uc?export=download&id=0Bz3fihKG133ceWNFQTQ5S0xhZUk"

_DATASET_FILENAME = "description2code_current.zip"

_DATASET_PB_PATH = "description2code_current/"

_DESC_DIR_NAME = "description"

_VOCAB_EN_FILENAME = "vocab.endefr"

_RE_CPP_INLINE_COMMENT = re.compile("//.*?\n") # Compiled once

# Constant defined for a language problem

CodingPbConstants = collections.namedtuple("CodingPbConstants", [

"code_dir_name",

"vocab_filename",

"filter_patterns",

"target_space",

])

PB_PY = CodingPbConstants(

code_dir_name="solutions_python",

vocab_filename="vocab.py",

filter_patterns=["#include", "# include", "import java."],

target_space=problem.SpaceID.PY_TOK,

)

PB_CPP = CodingPbConstants(

code_dir_name="solutions_c++",

vocab_filename="vocab.cpp",

filter_patterns=["import java."],

target_space=problem.SpaceID.CPP_TOK,

)

# Struct containing a coding problem (contains the paths to the descriptions

# and code files)

CodingPbInfo = collections.namedtuple("CodingPbInfo", "desc_file, code_files")

class Desc2CodeProblem(problem.Text2TextProblem):

"""Base class for Description2Code problems."""

@property

def is_character_level(self):

return False

@property

def num_shards(self):

return 10

@property

def use_subword_tokenizer(self):

return True

@property

def input_space_id(self):

return problem.SpaceID.EN_TOK

@property

def target_space_id(self):

return self.pb_constants.target_space

@property

def input_vocab_size(self):

return 2**15 # 32k

@property

def target_vocab_size(self):

return 2**12 # 4k

@property

def vocab_input_filename(self):

return "{}.{}".format(_VOCAB_EN_FILENAME, self.input_vocab_size)

@property

def vocab_target_filename(self):

return "{}.{}".format(

self.pb_constants.vocab_filename, self.target_vocab_size)

def preprocess_target(self, target):

"""Apply some preprocessing to the target.

For instance, remove space/tabs.

Args:

target (str): code source content

Returns:

the pre-processed string content

"""

return target

def feature_encoders(self, data_dir):

source_vocab_filename = os.path.join(data_dir, self.vocab_input_filename)

target_vocab_filename = os.path.join(data_dir, self.vocab_target_filename)

source_token = text_encoder.SubwordTextEncoder(source_vocab_filename)

target_token = text_encoder.SubwordTextEncoder(target_vocab_filename)

return {

"inputs": source_token,

"targets": target_token,

}

def generator(self, data_dir, tmp_dir, train):

# Called twice: for train and test

# Get the list of the training samples (coding challenge samples)

samples = list(generator_samples(tmp_dir, self.pb_constants))

# Split between train and dev

# Suffle to get problems from diverse sources (CodeChef and CodeForces) and

# dificulties in each set.

# Need to sort the samples first before shuffling (as walk() isn't

# deterministic)

samples.sort(key=lambda x: x.desc_file) # in-place

rng = random.Random(7531) # Local fixed seed

rng.shuffle(samples) # in-place

# Train: 5019/5228 problems

# Dev: 209/5228 problems

len_samples = len(samples)

split = len_samples // 25

samples = samples[split:] if train else samples[:split]

tf.logging.info("Number of samples for {}: {}/{}".format(

"train" if train else "dev",

len(samples),

len_samples

))

def generator_samples_content(get_source, get_target):

source, target = None, None

# Iterate over the coding samples

for sample in samples:

if get_source:

with tf.gfile.GFile(sample.desc_file, mode="r") as source_file:

source = source_file.read()

if get_target:

# Each challenge can have multiple implementations (or none)

for code_file in sample.code_files:

with tf.gfile.GFile(code_file, mode="r") as target_file:

target = target_file.read()

target = self.preprocess_target(target)

yield source, target

elif sample.code_files: # Only take the source if a target exists

yield source, target

def generator_target():

for _, target in generator_samples_content(False, True):

yield target.strip()

# Generate vocab for both source and target

source_vocab = generator_utils.get_or_generate_vocab(

data_dir, tmp_dir, self.vocab_input_filename, self.input_vocab_size)

target_vocab = generator_utils.get_or_generate_vocab_inner(

data_dir=data_dir,

vocab_filename=self.vocab_target_filename,

vocab_size=self.target_vocab_size,

generator=generator_target(),)

# Yield the training and testing samples

eos_list = [EOS]

for source, target in generator_samples_content(True, True):

source_ints = source_vocab.encode(source.strip()) + eos_list

target_ints = target_vocab.encode(target.strip()) + eos_list

yield {

"inputs": source_ints,

"targets": target_ints,

}

@registry.register_problem

class ProgrammingDesc2codePy(Desc2CodeProblem):

"""Description2Code for python problem."""

@property

def pb_constants(self):

return PB_PY

def preprocess_target(self, target):

"""Simple tab to space replacement."""

return target.replace("\t", " ")

@registry.register_problem

class ProgrammingDesc2codeCpp(Desc2CodeProblem):

"""Description2Code for C++ problem."""

@property

def pb_constants(self):

return PB_CPP

def preprocess_target(self, target):

"""Pre-process Cpp files."""

target = re.sub(_RE_CPP_INLINE_COMMENT, " ", target) # Remove comments

# The regex rule is quite simple, So will fail if a // is inside a string,

# and don't remove /* */ comments

target = " ".join(target.split()) # Normalize all spaces

return target

# Utils functions

def generator_samples(tmp_dir, pb_cst):

"""Generator for the dataset samples.

If not present, download and extract the dataset.

Args:

tmp_dir: path to the directory where to download the dataset.

pb_cst: CodingPbConstants object defining paths

Yields:

A CodingPbInfo object containing the next challenge informations.

"""

# Step1: Download dataset (eventually)

data_zip_path = generator_utils.maybe_download_from_drive(

directory=tmp_dir,

filename=_DATASET_FILENAME,

url=_DATASET_URL,

)

tf.logging.info("Data downloaded in: {}".format(data_zip_path))

# Step2: Extract dataset

# We could deduce _DATASET_PB_PATH from the zip file (instead of

# hardcoded path)

data_rootdir = os.path.join(tmp_dir, _DATASET_PB_PATH)

if not tf.gfile.Exists(data_rootdir):

with zipfile.ZipFile(data_zip_path, "r") as corpus_zip:

corpus_zip.extractall(tmp_dir)

# We could remove the extracted __MACOSX folder

tf.logging.info("Data extracted in: {}".format(tmp_dir))

else:

tf.logging.info("Data already extracted in: {}".format(tmp_dir))

# Step3: Extract the problems list on the extracted folder

def contains_samples(subdir, dirs, files): # pylint: disable=unused-argument

"""Check that the folder contains a problem."""

return (

_DESC_DIR_NAME in dirs and

pb_cst.code_dir_name in dirs

)

def next_sample(subdir, dirs, files): # pylint: disable=unused-argument

"""Return the filenames of the problem."""

# More could be extracted (like the expected inputs/outputs

# pairs, the problem difficulty, the names of the algorithmic techniques

# needed)

desc_file = os.path.join(subdir, _DESC_DIR_NAME, "description.txt")

code_files = []

# As the dataset is noisy, the program deduce the language from the file

# content.

code_pattern = os.path.join(subdir, pb_cst.code_dir_name, "*.txt")

for f in tf.gfile.Glob(code_pattern):

with tf.gfile.GFile(f, mode="r") as target_file:

# Hack to filter C++/Java files. In theory some python comments could

# make the file be concidered as C++ but in practice the chance of

# getting a false negative is low.

content = target_file.read()

if not any(p in content for p in pb_cst.filter_patterns):

code_files.append(f)

return CodingPbInfo(

desc_file=desc_file,

code_files=code_files

)

# The dataset contains problem from two different sources (CodeChef

# and CodeForces). Due to the limited number of samples, all problems from

# both sources are merged

for w in tf.gfile.Walk(data_rootdir):

if contains_samples(*w):

yield next_sample(*w)