<jupyter_start><jupyter_text>Install the library and download the pretrained models<jupyter_code>print("Installing dependencies...")
%tensorflow_version 2.x
!pip install -q t5==0.6.4

import functools
import os
import time
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds

import t5

!wget "https://www.dropbox.com/sh/kjoqdpj7e16dny9/AADdvjWVFckCgNQN-AqMKhiDa?dl=1" -O vocabulary.zip
!unzip vocabulary.zip
!rm vocabulary.zip
!wget "https://www.dropbox.com/sh/epth1k96g8ec4kc/AADiPvmSTTvGwxfbEt25binQa?dl=1" -O java.zip
!unzip java.zip
!rm java.zip<jupyter_output>Installing dependencies...
INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0 
--2020-11-09 16:01:12--  https://www.dropbox.com/sh/kjoqdpj7e16dny9/AADdvjWVFckCgNQN-AqMKhiDa?dl=1
Resolving www.dropbox.com (www.dropbox.com)... 172.16.17.32, fc00:db20:35b:7399::5:101
Connecting to www.dropbox.com (www.dropbox.com)|172.16.17.32|:443... connected.
HTTP request sent, awaiting response... 301 Moved Permanently
Location: /sh/dl/kjoqdpj7e16dny9/AADdvjWVFckCgNQN-AqMKhiDa [following]
--2020-11-09 16:01:13--  https://www.dropbox.com/sh/dl/kjoqdpj7e16dny9/AADdvjWVFckCgNQN-AqMKhiDa
Reusing existing connection to www.dropbox.com:443.
HTTP request sent, awaiting response... 302 Found
Location: https://uc20f4fe33865c423371518d45a5.dl.dropboxusercontent.com/zip_download_get/AmSGwCSdvXAFhCmCU7r2buBl1XutQ-t9d_vvEqIC9AJc3Qzx5t9caiQjUfvTLKDkPJ8NI2UckacyCAEovhcYbDOFW5EY-MltRDpbRtMhwZIesA?dl=1 [following]
--2020-11-09 16:01:13--  https://uc20f4fe[...]<jupyter_text>Set sentencepiece model<jupyter_code>from t5.data.sentencepiece_vocabulary import SentencePieceVocabulary

vocab_model_path = 'code_spm_unigram_40M.model'
vocab = SentencePieceVocabulary(vocab_model_path, extra_ids=100)

print("Vocab has a size of %d\n" % vocab.vocab_size)<jupyter_output>Vocab has a size of 32100<jupyter_text>Set the preprocessors and the task registry for the t5 model<jupyter_code>def java_codeSearchNet_dataset_fn(split, shuffle_files=False):
    del shuffle_files

    ds = tf.data.TextLineDataset(java_path[split])
    ds = ds.map(
        functools.partial(tf.io.decode_csv, record_defaults=["", ""], field_delim="\t", use_quote_delim=False),
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    ds = ds.map(lambda *ex: dict(zip(["code", "docstring"], ex)))
    return ds


def java_preprocessor(ds):
    def normalize_text(text):
        return text

    def to_inputs_and_targets(ex):
        return {
            "inputs": tf.strings.join(["function documentation generation java: ", normalize_text(ex["code"])]),
            "targets": normalize_text(ex["docstring"])
        }

    return ds.map(to_inputs_and_targets, num_parallel_calls=tf.data.experimental.AUTOTUNE)


t5.data.TaskRegistry.remove('function_documentation_generation_java_code')
t5.data.TaskRegistry.add(
    "function_documentation_generation_java_code",
    dataset_fn=java_codeSearchNet_dataset_fn,
    output_features={
        "inputs": t5.data.utils.Feature(vocabulary=vocab),
        "targets": t5.data.utils.Feature(vocabulary=vocab),
    },
    splits=["train", "validation"],
    text_preprocessor=[java_preprocessor],
    postprocess_fn=t5.data.postprocessors.lower_text,
    metric_fns=[t5.evaluation.metrics.bleu, t5.evaluation.metrics.accuracy, t5.evaluation.metrics.rouge],
)<jupyter_output><empty_output><jupyter_text>Set t5 small model<jupyter_code>MODEL_DIR = "small"
model_parallelism = 1
train_batch_size = 256

tf.io.gfile.makedirs(MODEL_DIR)
model = t5.models.MtfModel(
    model_dir=MODEL_DIR,
    tpu=None,
    tpu_topology=None,
    model_parallelism=model_parallelism,
    batch_size=train_batch_size,
    sequence_length={"inputs": 512, "targets": 512},
    mesh_shape="model:1,batch:1",
    mesh_devices=["GPU:0"],
    learning_rate_schedule=0.003,
    save_checkpoints_steps=5000,
    keep_checkpoint_max=None,
    iterations_per_loop=100,
)<jupyter_output><empty_output><jupyter_text>Code Documentation Summarization  Give the code for summarization<jupyter_code>code = """public static <T, U> Function<T, U> castFunction(Class<U> target) {\n        return new CastToClass<T, U>(target);\n    }""" #@param {type:"raw"}<jupyter_output><empty_output><jupyter_text>Parsing and Tokenization<jupyter_code>!pip install tree_sitter
!git clone https://github.com/tree-sitter/tree-sitter-java
from tree_sitter import Language, Parser

Language.build_library(
  'build/my-languages.so',
  ['tree-sitter-java']
)

JAVA_LANGUAGE = Language('build/my-languages.so', 'java')
parser = Parser()
parser.set_language(JAVA_LANGUAGE)
tree = parser.parse(bytes(code, "utf8"))
def my_traverse(node, code_list):
  lines = code.split('\n')
  if node.child_count == 0:
    line_start = node.start_point[0]
    line_end = node.end_point[0]
    char_start = node.start_point[1]
    char_end = node.end_point[1]
    if line_start != line_end:
        code_list.append(' '.join([lines[line_start][char_start:]] + lines[line_start+1:line_end] + [lines[line_end][:char_end]]))
    else:
        code_list.append(lines[line_start][char_start:char_end])
  else:
    for n in node.children:
      my_traverse(n, code_list)
  return ' '.join(code_list)
code_list=[]
tokenized_code = my_traverse(tree.root_node, code_list)
print("Output after tokenization: " + tokenized_code)<jupyter_output>Output after tokenization: public static < T , U > Function < T , U > castFunction  ( Class < U > target ) { return new CastToClass < T , U > ( target ) ; }<jupyter_text>Record the code for summarization with the prefix to a txt file<jupyter_code>codes = [tokenized_code]

inputs_path = 'input.txt'
with tf.io.gfile.GFile(inputs_path, "w") as f:
  for c in codes:
    f.write("function documentation generation java: %s\n" % c)

predict_outputs_path = 'MtfModel-output.txt'<jupyter_output><empty_output><jupyter_text>Running the model with the best checkpoint to summarize the given code<jupyter_code>model.batch_size = 8  # Min size for small model on v2-8 with parallelism 1.
model.predict(
    input_file="input.txt",
    output_file=predict_outputs_path,
    checkpoint_steps=60000,
    beam_size=4,
    vocabulary=vocab, 
    # Select the most probable output token at each step.
    temperature=0,
)<jupyter_output>INFO:tensorflow:Using config: {'_model_dir': 'small', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=None, num_cores_per_replica=1, per_host_input_for_training=4, tpu_job_name=None, initial_infeed_sleep_s[...]<jupyter_text>Code Summarization Result<jupyter_code>prediction_file = "MtfModel-output.txt-60000"
print("\nPredictions using checkpoint 60000:\n" )
with tf.io.gfile.GFile(prediction_file) as f:
  for c, d in zip(codes, f):
    if c:
      print("Code for prediction: " + c + '\n')
      print("Generated Documentation: " + d)<jupyter_output>Predictions using checkpoint 60000:

Code for prediction: public static < T , U > Function < T , U > castFunction  ( Class < U > target ) { return new CastToClass < T , U > ( target ) ; }

Generated Documentation: b'Casts a Cast to another type'
{'fraction_non_alphanumeric': 0.12371778834125594, 'fraction_numerical': 0.03365023767825869, 'mean_word_length': 5.6625, 'pattern_counts': {'":': 6, '<': 44, '<?xml version=': 0, '>': 44, 'https://': 7, 'lorem ipsum': 0, 'www.': 9, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '84316', 'license': 'MIT', 'n_tokens_mistral': 3126, 'n_tokens_neox': 2926, 'n_words': 518, 'path': 'prediction/single task/function documentation generation/java/t5 interface/small_model.ipynb', 'repo_name': 'victory-hash/CodeTrans', 'size': '7992'}