Cleanup

app.py

@@ -3,28 +3,33 @@ import requests
 from transformers import AutoTokenizer, T5ForConditionalGeneration, AutoModelForSeq2SeqLM
 import torch
 
-tokenizer = AutoTokenizer.from_pretrained("microsoft/codereviewer")
-
-tokenizer.special_dict = {
-    f"<e{i}>": tokenizer.get_vocab()[f"<e{i}>"] for i in range(99, -1, -1)
-}
-tokenizer.mask_id = tokenizer.get_vocab()["<mask>"]
-tokenizer.bos_id = tokenizer.get_vocab()["<s>"]
-tokenizer.pad_id = tokenizer.get_vocab()["<pad>"]
-tokenizer.eos_id = tokenizer.get_vocab()["</s>"]
-tokenizer.msg_id = tokenizer.get_vocab()["<msg>"]
-tokenizer.keep_id = tokenizer.get_vocab()["<keep>"]
-tokenizer.add_id = tokenizer.get_vocab()["<add>"]
-tokenizer.del_id = tokenizer.get_vocab()["<del>"]
-tokenizer.start_id = tokenizer.get_vocab()["<start>"]
-tokenizer.end_id = tokenizer.get_vocab()["<end>"]
-
-model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/codereviewer")
-
-model.eval()
 
 MAX_SOURCE_LENGTH = 512
 
+
+def prepare_models():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/codereviewer")
+
+    tokenizer.special_dict = {
+        f"<e{i}>": tokenizer.get_vocab()[f"<e{i}>"] for i in range(99, -1, -1)
+    }
+    tokenizer.mask_id = tokenizer.get_vocab()["<mask>"]
+    tokenizer.bos_id = tokenizer.get_vocab()["<s>"]
+    tokenizer.pad_id = tokenizer.get_vocab()["<pad>"]
+    tokenizer.eos_id = tokenizer.get_vocab()["</s>"]
+    tokenizer.msg_id = tokenizer.get_vocab()["<msg>"]
+    tokenizer.keep_id = tokenizer.get_vocab()["<keep>"]
+    tokenizer.add_id = tokenizer.get_vocab()["<add>"]
+    tokenizer.del_id = tokenizer.get_vocab()["<del>"]
+    tokenizer.start_id = tokenizer.get_vocab()["<start>"]
+    tokenizer.end_id = tokenizer.get_vocab()["<end>"]
+
+    model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/codereviewer")
+
+    model.eval()
+    return tokenizer, model
+
+
 def pad_assert(tokenizer, source_ids):
     source_ids = source_ids[:MAX_SOURCE_LENGTH - 2]
     source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
@@ -73,19 +78,24 @@ class FileDiffs(object):
             self.diffs[-1] += "\n" + line
 
 
-def review_commit(user, repository, commit):
+def review_commit(user="p4vv37", repository="ueflow", commit="610a8c7b02b946bc9e5e26e6dacbba0e2abba259"):
+    tokenizer, model = prepare_models()
+
+    # Get diff and commit metadata from GitHub API
     commit_metadata = requests.get(F"https://api.github.com/repos/{user}/{repository}/commits/{commit}").json()
     msg = commit_metadata["commit"]["message"]
-
-    diff_data = requests.get(F"https://api.github.com/repos/{user}/{repository}/commits/{commit}", headers={"Accept":"application/vnd.github.diff"})
+    diff_data = requests.get(F"https://api.github.com/repos/{user}/{repository}/commits/{commit}",
+                             headers={"Accept": "application/vnd.github.diff"})
     code_diff = diff_data.text
 
+    # Parse diff into FileDiffs objects
     files_diffs = list()
     for file in code_diff.split("diff --git"):
         if len(file) > 0:
             fd = FileDiffs(file)
             files_diffs.append(fd)
 
+    # Generate comments for each diff
     output = ""
     for fd in files_diffs:
         output += F"File:{fd.file_path}\n"
@@ -103,11 +113,22 @@ def review_commit(user, repository, commit):
                                    num_return_sequences=2
                                    )
             preds = list(preds.cpu().numpy())
-            pred_nls = [tokenizer.decode(id[2:], skip_special_tokens=True, clean_up_tokenization_spaces=False) for id in
-                        preds]
+            pred_nls = [tokenizer.decode(_id[2:], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+                        for _id in preds]
             output += diff + "\n#######\nComment:\n#######\n" + pred_nls[0] + "\n#######\n"
     return output
 
 
-iface = gr.Interface(fn=review_commit, inputs=["text", "text", "text"], outputs="text")
+description = "An interface for running " \
+              "\"Microsoft CodeBERT CodeReviewer: Pre-Training for Automating Code Review Activities.\" " \
+              "(microsoft/codereviewer) on GitHub commits."
+examples = [
+    ["p4vv37", "ueflow", "610a8c7b02b946bc9e5e26e6dacbba0e2abba259"],
+    ["microsoft", "vscode", "378b0d711f6b82ac59b47fb246906043a6fb995a"],
+]
+iface = gr.Interface(fn=review_commit,
+                     description=description,
+                     inputs=["text", "text", "text"],
+                     outputs="text",
+                     examples=examples)
 iface.launch()