Deploy to HF Space

clear.py

@@ -0,0 +1,49 @@
+from clearml import Model
+import torch
+import os
+# Import needed classes for local loading and LoRA construction
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+
+# 1. Download the LoRA checkpoint artifact from ClearML
+CLEARML_MODEL_ID = "34e25deb24c64b74b29c8519ed15fe3e"
+model_obj = Model(model_id=CLEARML_MODEL_ID)
+checkpoint_path = model_obj.get_local_copy()
+adapter_dir = os.path.dirname(checkpoint_path)
+print(f"LoRA checkpoint downloaded to: {checkpoint_path}")
+
+# 2. Load the base pretrained CodeT5 model and tokenizer from local config.json directory
+BASE_MODEL_PATH = "microsoft/codereviewer"
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH)
+base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL_PATH)
+
+# Print all base model parameters and their shapes
+print("\nBase model parameters:")
+for name, param in base_model.named_parameters():
+    print(f"{name}: {tuple(param.shape)}")
+
+# 3. Reconstruct and attach LoRA adapters
+lora_config = LoraConfig(
+    r=64,
+    lora_alpha=128,
+    target_modules=["q", "k", "v", "o", "wi", "wo"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="SEQ_2_SEQ_LM"
+)
+model = get_peft_model(base_model, lora_config)
+
+# 4. Load LoRA adapter weights from ClearML checkpoint
+adapter_state = torch.load(checkpoint_path, map_location="cpu")
+model.load_state_dict(adapter_state, strict=False)
+
+# 5. Move to CPU and set evaluation mode
+model.to("cpu").eval()
+
+print("Model with LoRA adapters loaded and ready for inference.")
+
+# Print out all LoRA adapter parameter names and shapes as before
+print("\nFinetuned (LoRA adapter) parameters:")
+for name, param in model.named_parameters():
+    if "lora_" in name:
+        print(f"{name}: {tuple(param.shape)}")

configs.py

@@ -0,0 +1,252 @@
+import random
+import torch
+import logging
+import multiprocessing
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def add_args(parser):
+    parser.add_argument(
+        "--task",
+        type=str,
+        required=False,
+        choices=[
+            "review",
+        ],
+    )
+    parser.add_argument(
+        "--model_type",
+        default="codet5",
+        type=str,
+        choices=["roberta", "t5", "bart", "codet5", "scratch"],
+    )
+    parser.add_argument("--add_lang_ids", action="store_true")
+    parser.add_argument("--from_scratch", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--start_epoch", default=0, type=int)
+    parser.add_argument("--train_epochs", default=10, type=int)
+    parser.add_argument("--tokenizer_path", type=str, required=False)
+    
+    parser.add_argument(
+        "--output_dir",
+        default=None,
+        type=str,
+        required=False,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--load_model_path",
+        default=None,
+        type=str,
+        required=False
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        default=None,
+        type=str,
+        help="Path to trained model: Should contain the .bin files",
+    )
+    ## Other parameters
+    parser.add_argument(
+        "--train_path",
+        default=None,
+        type=str,
+        help="The pretrain files path. Should contain the .jsonl files for this task.",
+    )
+    parser.add_argument(
+        "--eval_chunkname",
+        default=None,
+        type=str,
+        help="The eval file name.",
+    )
+    parser.add_argument(
+        "--train_filename",
+        default=None,
+        type=str,
+        help="The train filename. Should contain the .jsonl files for this task.",
+    )
+    parser.add_argument(
+        "--dev_filename",
+        default=None,
+        type=str,
+        help="The dev filename. Should contain the .jsonl files for this task.",
+    )
+    parser.add_argument(
+        "--test_filename",
+        default=None,
+        type=str,
+        help="The test filename. Should contain the .jsonl files for this task.",
+    )
+    parser.add_argument(
+        "--gold_filename",
+        default=None,
+        type=str,
+        help="The gold filename. Should contain the .jsonl files for this task.",
+    )
+    parser.add_argument(
+        "--config_name",
+        default="Salesforce/codet5-base",
+        type=str,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--max_source_length",
+        default=64,
+        type=int,
+        help="The maximum total source sequence length after tokenization. Sequences longer "
+        "than this will be truncated, sequences shorter will be padded.",
+    )
+    parser.add_argument(
+        "--max_target_length",
+        default=32,
+        type=int,
+        help="The maximum total target sequence length after tokenization. Sequences longer "
+        "than this will be truncated, sequences shorter will be padded.",
+    )
+    parser.add_argument(
+        "--do_train", action="store_true", help="Whether to run eval on the train set."
+    )
+    parser.add_argument(
+        "--do_eval", action="store_true", help="Whether to run eval on the dev set."
+    )
+    parser.add_argument(
+        "--do_test", action="store_true", help="Whether to run eval on the dev set."
+    )
+    parser.add_argument(
+        "--raw_input", action="store_true", help="Whether to use simple input format (set for baselines)."
+    )
+    parser.add_argument(
+        "--do_lower_case",
+        action="store_true",
+        help="Set this flag if you are using an uncased model.",
+    )
+    parser.add_argument(
+        "--no_cuda", action="store_true", help="Avoid using CUDA when available"
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        default=8,
+        type=int,
+        help="Batch size per GPU/CPU for training.",
+    )
+    parser.add_argument(
+        "--eval_batch_size",
+        default=8,
+        type=int,
+        help="Batch size per GPU/CPU for evaluation.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        default=5e-5,
+        type=float,
+        help="The initial learning rate for Adam.",
+    )
+    parser.add_argument(
+        "--mask_rate", default=0.15, type=float, help="The masked rate of input lines.",
+    )
+    parser.add_argument(
+        "--beam_size", default=6, type=int, help="beam size for beam search"
+    )
+    parser.add_argument(
+        "--weight_decay", default=0.0, type=float, help="Weight deay if we apply some."
+    )
+    parser.add_argument(
+        "--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer."
+    )
+    parser.add_argument(
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
+    )
+    parser.add_argument(
+        "--save_steps", default=-1, type=int,
+    )
+    parser.add_argument(
+        "--log_steps", default=-1, type=int,
+    )
+    parser.add_argument("--eval_steps", default=-1, type=int, help="")
+    parser.add_argument("--eval_file", default="", type=str)
+    parser.add_argument("--out_file", default="", type=str)
+    parser.add_argument("--break_cnt", default=-1, type=int)
+    parser.add_argument("--train_steps", default=-1, type=int, help="")
+    parser.add_argument(
+        "--warmup_steps", default=100, type=int, help="Linear warmup over warmup_steps."
+    )
+    parser.add_argument(
+        "--gpu_per_node",
+        type=int,
+        default=4,
+        help="gpus per node",
+    )
+    parser.add_argument(
+        "--node_index",
+        type=int,
+        default=0,
+        help="For distributed training: node_index",
+    )
+    parser.add_argument(
+        "--local_rank",
+        type=int,
+        default=-1,
+        help="For distributed training: local_rank",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=2233, help="random seed for initialization"
+    )  # previous one 42
+    # Or in configs.py if add_args is defined there
+
+    parser.add_argument(
+        "--clearml_train_dataset_id",
+        type=str,
+        default=None,
+        help="ClearML Dataset ID to fetch training data from. Overrides train_filename if provided.",
+    )
+    parser.add_argument(
+        "--clearml_valid_dataset_id",
+        type=str,
+        default=None,
+        help="ClearML Dataset ID to fetch validation data from. Overrides dev_filename if provided.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def set_dist(args):
+    # Setup CUDA, GPU & distributed training
+    if args.local_rank == -1 or args.no_cuda:
+        device = torch.device(
+            "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"
+        )
+        args.n_gpu = torch.cuda.device_count()
+    else:
+        # Setup for distributed data parallel
+        torch.cuda.set_device(args.local_rank)
+        device = torch.device("cuda", args.local_rank)
+        torch.distributed.init_process_group(backend="nccl")
+        args.n_gpu = 1
+    cpu_count = multiprocessing.cpu_count()
+    logger.warning(
+        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, cpu count: %d",
+        args.local_rank,
+        device,
+        args.n_gpu,
+        bool(args.local_rank != -1),
+        cpu_count,
+    )
+    args.device = device
+    args.cpu_count = cpu_count
+
+
+def set_seed(args):
+    """set random seed."""
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    # if args.n_gpu > 0:
+    torch.cuda.manual_seed_all(args.seed)

extract_pr_comment.py

@@ -0,0 +1,46 @@
+import os
+import json
+import base64
+import requests
+from github import Github
+
+# Path to the GitHub Actions event payload
+event_path = os.environ.get("GITHUB_EVENT_PATH")
+if not event_path or not os.path.exists(event_path):
+    print("No event payload found.")
+    exit(1)
+
+with open(event_path, "r") as f:
+    event = json.load(f)
+
+# Only proceed if this is a PR comment event
+if "pull_request" not in event.get("issue", {}):
+    print("Not a PR comment event.")
+    exit(0)
+
+pr_number = event["issue"]["number"]
+comment_body = event["comment"]["body"]
+repo_full_name = event["repository"]["full_name"]
+token = os.environ.get("GITHUB_TOKEN")
+
+if not token:
+    print("No GITHUB_TOKEN found in environment.")
+    exit(1)
+
+gh = Github(token)
+repo = gh.get_repo(repo_full_name)
+pr = repo.get_pull(pr_number)
+
+files = []
+for file in pr.get_files():
+    cf = repo.get_contents(file.filename, ref=pr.head.sha)
+    content = base64.b64decode(cf.content).decode("utf-8")
+    files.append({"filename": file.filename, "content": content})
+
+fastapi_url = "http://127.0.0.1:8000/pr-comments"
+payload = {
+    "comment": comment_body,
+    "files": files
+}
+response = requests.post(fastapi_url, json=payload)
+print(f"FastAPI response: {response.status_code} {response.text}")

fastapi_app.py

@@ -0,0 +1,153 @@
+from fastapi import FastAPI, Request, Form
+from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
+from pydantic import BaseModel
+from typing import List
+from clearml import Model
+import torch
+from configs import add_args
+from models import build_or_load_gen_model
+import argparse
+from argparse import Namespace
+import os
+from peft import PeftModel, PeftConfig, get_peft_model, LoraConfig
+
+MAX_SOURCE_LENGTH = 512
+
+def pad_assert(tokenizer, source_ids):
+    source_ids = source_ids[:MAX_SOURCE_LENGTH - 2]
+    source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
+    pad_len = MAX_SOURCE_LENGTH - len(source_ids)
+    source_ids += [tokenizer.pad_id] * pad_len
+    assert len(source_ids) == MAX_SOURCE_LENGTH, "Not equal length."
+    return source_ids
+
+# Encode code content and comment into model input
+def encode_diff(tokenizer, code, comment):
+    # Tokenize code file content
+    code_ids = tokenizer.encode(code, max_length=MAX_SOURCE_LENGTH, truncation=True)[1:-1]
+    # Tokenize comment
+    comment_ids = tokenizer.encode(comment, max_length=MAX_SOURCE_LENGTH, truncation=True)[1:-1]
+    # Concatenate: [BOS] + code + [EOS] + [msg_id] + comment
+    source_ids = [tokenizer.bos_id] + code_ids + [tokenizer.eos_id]
+    source_ids += [tokenizer.msg_id] + comment_ids
+    # Pad/truncate to fixed length
+    source_ids = source_ids[:MAX_SOURCE_LENGTH - 2]
+    source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
+    pad_len = MAX_SOURCE_LENGTH - len(source_ids)
+    source_ids += [tokenizer.pad_id] * pad_len
+    assert len(source_ids) == MAX_SOURCE_LENGTH, "Not equal length."
+    return source_ids
+
+# Load base model architecture and tokenizer from HuggingFace
+BASE_MODEL_NAME = "microsoft/codereviewer"
+args = Namespace(
+    model_name_or_path=BASE_MODEL_NAME,
+    load_model_path=None,
+    # Add other necessary default arguments if build_or_load_gen_model requires them
+)
+print(f"Loading base model architecture and tokenizer from: {BASE_MODEL_NAME}")
+config, base_model, tokenizer = build_or_load_gen_model(args)
+print("Base model architecture and tokenizer loaded.")
+
+# Download the fine-tuned weights from ClearML
+CLEARML_MODEL_ID = "34e25deb24c64b74b29c8519ed15fe3e"
+model_obj = Model(model_id=CLEARML_MODEL_ID)
+finetuned_weights_path = model_obj.get_local_copy()
+adapter_dir = os.path.dirname(finetuned_weights_path)
+
+print(f"Fine-tuned adapter weights downloaded to directory: {adapter_dir}")
+
+# Create LoRA configuration matching the fine-tuned checkpoint
+lora_cfg = LoraConfig(
+    r=64,
+    lora_alpha=128,
+    target_modules=["q", "wo", "wi", "v", "o", "k"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="SEQ_2_SEQ_LM"
+)
+# Wrap base model with PEFT LoRA
+peft_model = get_peft_model(base_model, lora_cfg)
+# Load adapter-only weights and merge into base
+adapter_state = torch.load(finetuned_weights_path, map_location="cpu")
+peft_model.load_state_dict(adapter_state, strict=False)
+model = peft_model.merge_and_unload()
+print("Merged base model with LoRA adapters.")
+
+model.to("cpu")
+model.eval()
+print("Model ready for inference.")
+
+app = FastAPI()
+
+last_payload = {"comment": "", "files": []}
+last_infer_result = {"generated_code": ""}
+
+class FileContent(BaseModel):
+    filename: str
+    content: str
+
+class PRPayload(BaseModel):
+    comment: str
+    files: List[FileContent]
+
+class InferenceRequest(BaseModel):
+    comment: str
+    files: List[FileContent]
+
+
+@app.get("/")
+def root():
+    return {"message": "FastAPI PR comment service is running"}
+
+@app.post("/pr-comments")
+async def receive_pr_comment(payload: PRPayload):
+    global last_payload
+    last_payload = payload.dict()
+    # Return the received payload as JSON and also redirect to /show
+    return JSONResponse(content={"status": "received", "payload": last_payload, "redirect": "/show"})
+
+@app.get("/show", response_class=HTMLResponse)
+def show_last_comment():
+    html = f"<h2>Received Comment</h2><p>{last_payload['comment']}</p><hr>"
+    for file in last_payload["files"]:
+        html += f"<h3>{file['filename']}</h3><pre>{file['content']}</pre><hr>"
+    return html
+
+@app.post("/infer")
+async def infer(request: InferenceRequest):
+    global last_infer_result
+    print("[DEBUG] Received /infer request with:", request.dict())
+    
+    code = request.files[0].content if request.files else ""
+    source_ids = encode_diff(tokenizer, code, request.comment)
+    # print("[DEBUG] source_ids:", source_ids)
+    #tokens = [tokenizer.decode([sid], skip_special_tokens=False) for sid in source_ids]
+    #print("[DEBUG] tokens:", tokens)
+    inputs = torch.tensor([source_ids], dtype=torch.long)
+    inputs_mask = inputs.ne(tokenizer.pad_id)
+    
+    preds = model.generate(
+        inputs,
+        attention_mask=inputs_mask,
+        use_cache=True,
+        num_beams=5,
+        early_stopping=True,
+        max_length=100,
+        num_return_sequences=1
+    )
+    
+    pred = preds[0].cpu().numpy()
+    pred_nl = tokenizer.decode(pred[2:], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    last_infer_result = {"generated_code": pred_nl}
+    return last_infer_result
+
+@app.get("/show-infer", response_class=HTMLResponse)
+def show_infer_result():
+    html = f"<h2>Generated Message</h2><pre>{last_infer_result['generated_code']}</pre>"
+    return html
+
+if __name__ == "__main__":
+    # Place any CLI/training logic here if needed
+    # This block is NOT executed when running with uvicorn
+    pass

models.py

@@ -0,0 +1,208 @@
+import os
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss, BCEWithLogitsLoss
+import numpy as np
+from utils import MyTokenizer
+from transformers import (
+    RobertaConfig,
+    RobertaModel,
+    RobertaTokenizer,
+    BartConfig,
+    BartForConditionalGeneration,
+    BartTokenizer,
+    T5Config,
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ReviewerModel(T5ForConditionalGeneration):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.cls_head = nn.Linear(self.config.d_model, 2, bias=True)
+        self.init()
+
+    def init(self):
+        nn.init.xavier_uniform_(self.lm_head.weight)
+        factor = self.config.initializer_factor
+        self.cls_head.weight.data.normal_(mean=0.0, \
+            std=factor * ((self.config.d_model) ** -0.5))
+        self.cls_head.bias.data.zero_()
+
+    def forward(
+        self, *argv, **kwargs
+    ):
+        r"""
+        Doc from Huggingface transformers:
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[-100, 0, ...,
+            config.vocab_size - 1]`. All labels set to ``-100`` are ignored (masked), the loss is only computed for
+            labels in ``[0, ..., config.vocab_size]``
+        Returns:
+        Examples::
+            >>> from transformers import T5Tokenizer, T5ForConditionalGeneration
+            >>> tokenizer = T5Tokenizer.from_pretrained('t5-small')
+            >>> model = T5ForConditionalGeneration.from_pretrained('t5-small')
+            >>> # training
+            >>> input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids
+            >>> labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2>', return_tensors='pt').input_ids
+            >>> outputs = model(input_ids=input_ids, labels=labels)
+            >>> loss = outputs.loss
+            >>> logits = outputs.logits
+            >>> # inference
+            >>> input_ids = tokenizer("summarize: studies have shown that owning a dog is good for you", return_tensors="pt").input_ids  # Batch size 1
+            >>> outputs = model.generate(input_ids)
+            >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+            >>> # studies have shown that owning a dog is good for you.
+        """
+        if "cls" in kwargs:
+            assert (
+                "input_ids" in kwargs and \
+                "labels" in kwargs and \
+                "attention_mask" in kwargs
+            )
+            return self.cls(
+                input_ids=kwargs["input_ids"],
+                labels=kwargs["labels"],
+                attention_mask=kwargs["attention_mask"],
+            )
+        if "input_labels" in kwargs:
+            assert (
+                "input_ids" in kwargs and \
+                "input_labels" in kwargs and \
+                "decoder_input_ids" in kwargs and \
+                "attention_mask" in kwargs and \
+                "decoder_attention_mask" in kwargs
+            ), "Please give these arg keys."
+            input_ids = kwargs["input_ids"]
+            input_labels = kwargs["input_labels"]
+            decoder_input_ids = kwargs["decoder_input_ids"]
+            attention_mask = kwargs["attention_mask"]
+            decoder_attention_mask = kwargs["decoder_attention_mask"]
+            if "encoder_loss" not in kwargs:
+                encoder_loss = True
+            else:
+                encoder_loss = kwargs["encoder_loss"]
+            return self.review_forward(input_ids, input_labels, decoder_input_ids, attention_mask, decoder_attention_mask, encoder_loss)
+        return super().forward(*argv, **kwargs)
+
+    def cls(
+        self,
+        input_ids,
+        labels,
+        attention_mask,
+    ):
+        encoder_outputs = self.encoder( \
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            return_dict=False
+        )
+        hidden_states = encoder_outputs[0]
+        first_hidden = hidden_states[:, 0, :]
+        first_hidden = nn.Dropout(0.3)(first_hidden)
+        logits = self.cls_head(first_hidden)
+        loss_fct = CrossEntropyLoss()
+        if labels != None:
+            loss = loss_fct(logits, labels)
+            return loss
+        return logits
+
+    def review_forward(
+        self,
+        input_ids,
+        input_labels,
+        decoder_input_ids,
+        attention_mask,
+        decoder_attention_mask,
+        encoder_loss=True
+    ):
+        encoder_outputs = self.encoder( \
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            return_dict=False
+        )
+        hidden_states = encoder_outputs[0]
+        decoder_inputs = self._shift_right(decoder_input_ids)
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_inputs,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            output_attentions=False,
+            return_dict=False
+        )
+        sequence_output = decoder_outputs[0]
+        if self.config.tie_word_embeddings: # this is True default
+            sequence_output = sequence_output * (self.model_dim ** -0.5)
+        if encoder_loss:
+            # print(self.encoder.get_input_embeddings().weight.shape)
+            cls_logits = nn.functional.linear(hidden_states, self.encoder.get_input_embeddings().weight)
+            # cls_logits = self.cls_head(hidden_states)
+        lm_logits = self.lm_head(sequence_output)
+        if decoder_input_ids is not None:
+            lm_loss_fct = CrossEntropyLoss(ignore_index=0)      # Warning: PAD_ID should be 0
+            loss = lm_loss_fct(lm_logits.view(-1, lm_logits.size(-1)), decoder_input_ids.view(-1))
+            if encoder_loss and input_labels is not None:
+                cls_loss_fct = CrossEntropyLoss(ignore_index=-100)
+                loss += cls_loss_fct(cls_logits.view(-1, cls_logits.size(-1)), input_labels.view(-1))
+            return loss
+        return cls_logits, lm_logits
+
+def get_model_size(model):
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    model_size = sum([np.prod(p.size()) for p in model_parameters])
+    return "{}M".format(round(model_size / 1e6))
+
+
+def build_or_load_gen_model(args):
+    config_class, model_class, tokenizer_class = T5Config, ReviewerModel, RobertaTokenizer
+    
+    config = config_class.from_pretrained(args.model_name_or_path)
+    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
+    model = model_class.from_pretrained(args.model_name_or_path, config=config)
+
+    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>"]
+
+    logger.info(
+        "Finish loading model [%s] from %s",
+        get_model_size(model),
+        args.model_name_or_path,
+    )
+
+    if args.load_model_path is not None:
+        model_path = os.path.join(args.load_model_path, "pytorch_model.bin")
+        logger.info("Reload model from {}".format(model_path))
+        try:
+            model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        except RuntimeError:
+            saved = model.cls_head
+            model.cls_head = None
+            model.load_state_dict(torch.load(model_path, map_location="cpu"))
+            model.cls_head = saved
+        model.to(args.local_rank)
+
+    return config, model, tokenizer
+
+

test.py

@@ -0,0 +1,29 @@
+import base64          # For decoding Base64 content
+import requests        # For HTTP GET on raw_url
+from github import Github  # PyGithub
+
+# ==== CHANGE VALUES BELOW ====================================
+TOKEN     = "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"       # <-- Change: Your GitHub PAT
+OWNER     = "Habil7"      # <-- Change: Repo owner
+REPO_NAME = "git-demo"            # <-- Change: Repo name
+PR_NUMBER = 4                       # <-- Change: Pull request number
+# =============================================================
+
+gh   = Github(TOKEN)
+repo = gh.get_repo(f"{OWNER}/{REPO_NAME}")
+pr   = repo.get_pull(PR_NUMBER)
+print(pr)
+
+# Print PR comments
+print("\n--- PR Comments ---")
+for comment in pr.get_issue_comments():
+    print(f"{comment.user.login}: {comment.body}")
+
+print(f"Number of files in PR: {pr.get_files().totalCount}")
+
+for file in pr.get_files():
+    print(f"\n=== {file.filename} ===")
+    # Fetch and decode via PyGithub get_contents
+    cf = repo.get_contents(file.filename, ref=pr.head.sha)
+    content_via_api = base64.b64decode(cf.content).decode("utf-8")
+    print(content_via_api)

utils.py

@@ -0,0 +1,823 @@
+import re, json
+import os, random
+import torch, logging
+from copy import deepcopy as cp
+from torch.utils.data import Dataset
+from tokenizers import ByteLevelBPETokenizer
+from transformers import T5Tokenizer, RobertaTokenizer
+import nltk
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.INFO,
+)
+logger = logging.getLogger(__name__)
+
+
+
+class MyTokenizer(object):
+    """
+    Wrapper for ByteLevelBPETokenizer
+    """
+    def __init__(self, vocab=None, merges=None, **kwargs):
+        self.tokenizer = ByteLevelBPETokenizer(vocab, merges, **kwargs)
+        self.update_id2token()
+
+    @staticmethod
+    def from_pretrained(path):
+        vocabp = os.path.join(path, "vocab.json")
+        mergesp = os.path.join(path, "merges.txt")
+        mytoken = MyTokenizer(vocabp, mergesp)
+        return mytoken
+
+    def update_id2token(self):
+        vocab = self.tokenizer.get_vocab()
+        self.id2token = {vocab[token]: token for token in vocab}
+
+    def add_special_tokens(self, dic):
+        for values in dic.values():
+            self.tokenizer.add_special_tokens(values)
+        self.update_id2token()
+
+    def convert_ids_to_tokens(self, ids):
+        vocab = self.id2token
+        return [vocab[i] for i in ids]
+    
+    def decode(self, ids, **kwargs):    ##### to be update
+        tokens = self.convert_ids_to_tokens(ids)
+        return " ".join(tokens)
+
+    def encode(self, text, **kwargs):
+        text = text.encode("ascii", errors="ignore").decode("ascii")
+        return self.tokenizer.encode(text).ids
+
+    def get_vocab(self):
+        return self.tokenizer.get_vocab()
+
+    def __len__(self):
+        return len(self.tokenizer.get_vocab())
+
+
+class RefineFeatures(object):
+    def __init__(self, example_id, source_ids, target_ids):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids
+
+class RefineDataset(Dataset):
+    def __init__(self, tokenizer, pool, args, file_path, samplenum=-1):
+        self.tokenizer = tokenizer
+        self.args = args
+        logger.info("Reading examples from {}".format(file_path))
+        examples = [json.loads(line) for line in open(file_path)]
+        for i in range(len(examples)):
+            if "id" not in examples[i]:
+                examples[i]["id"] = i
+        if samplenum > 0:
+            examples = examples[:samplenum]
+        logger.info(f"Tokenize examples: {file_path}")
+        self.feats = pool.map(self.tokenize, \
+            [(example, tokenizer, args) for example in examples])
+        
+    def tokenize(self, item):
+        example, tokenizer, args = item
+        oldlines = example["old"].split("\n")
+        newlines = example["new"].split("\n")
+        oldlines = [line[1:].strip() for line in oldlines]
+        newlines = [line[1:].strip() for line in newlines]
+        oldlines = "\n".join(oldlines)
+        newlines = "\n".join(newlines)
+        oldlines = "<add>" + oldlines.replace("\n", "<add>")
+        newlines = "<add>" + newlines.replace("\n", "<add>")
+        comment = example["comment"]
+        srcids = self.encode_remove(tokenizer, oldlines, args)
+        srcids += [tokenizer.msg_id] + self.encode_remove(tokenizer, comment, args)
+        tgtids = self.encode_remove(tokenizer, newlines, args)
+        srcids, tgtids = self.pad_assert(srcids, tgtids, args, tokenizer)
+        return RefineFeatures(example["id"], srcids, tgtids)
+
+    @staticmethod
+    def process_pred_gold(pred, gold):
+        gold = gold.split("\n")
+        gold = [line[1:].strip() for line in gold]
+        gold = " ".join(gold)
+        pred = " ".join(pred.split())
+        pred = pred.replace("<add> ", "")
+        return pred, gold
+
+    def pad_assert(self, source_ids, target_ids, args, tokenizer):
+        source_ids = source_ids[:args.max_source_length - 2]
+        source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
+        pad_len = args.max_source_length - len(source_ids)
+        source_ids += [tokenizer.pad_id] * pad_len
+        target_ids = target_ids[:args.max_target_length - 2]
+        target_ids = [tokenizer.bos_id] + target_ids + [tokenizer.eos_id]
+        pad_len = args.max_target_length - len(target_ids)
+        target_ids += [tokenizer.pad_id] * pad_len
+        assert len(source_ids) == args.max_source_length, "Not equal length."
+        assert len(target_ids) == args.max_target_length, "Not equal length."
+        return source_ids, target_ids
+
+    def encode_remove(self, tokenizer, text, args):
+        text = tokenizer.encode(text, max_length=args.max_source_length, truncation=True)
+        if type(tokenizer) == T5Tokenizer:
+            return text[:-1]
+        elif type(tokenizer) == RobertaTokenizer:
+            return text[1:-1]
+        elif type(tokenizer) == MyTokenizer:
+            return text
+        else:
+            raise NotImplementedError
+
+    def __len__(self):
+        return len(self.feats)
+
+    def __getitem__(self, i):
+        return self.feats[i]
+
+class SimpleRefineDataset(RefineDataset):
+    def tokenize(self, item):
+        example, tokenizer, args = item
+        oldlines = example["old"].split("\n")
+        newlines = example["new"].split("\n")
+        oldlines = [line[1:].strip() for line in oldlines]
+        newlines = [line[1:].strip() for line in newlines]
+        oldlines = " ".join(oldlines)
+        newlines = " ".join(newlines)
+        comment = example["comment"]
+        srcids = self.encode_remove(tokenizer, oldlines, args)
+        srcids += [tokenizer.msg_id] + self.encode_remove(tokenizer, comment, args)
+        tgtids = self.encode_remove(tokenizer, newlines, args)
+        srcids, tgtids = self.pad_assert(srcids, tgtids, args, tokenizer)
+        return RefineFeatures(example["id"], srcids, tgtids)
+
+    @staticmethod
+    def process_pred_gold(pred, gold):
+        gold = gold.split("\n")
+        gold = [line[1:].strip() for line in gold]
+        gold = " ".join(gold)
+        pred = " ".join(pred.split())
+        return pred, gold
+
+
+class Seq2SeqDataset(RefineDataset):
+    def tokenize(self, item):
+        example, tokenizer, args = item
+        inputs, outputs = example["old"], example["new"]
+        inputs = " ".join(inputs.split())
+        outputs = " ".join(outputs.split())
+        srcids = self.encode_remove(tokenizer, inputs, args)
+        tgtids = self.encode_remove(tokenizer, outputs, args)
+        srcids, tgtids = self.pad_assert(srcids, tgtids, args, tokenizer)
+        return RefineFeatures(example["id"], srcids, tgtids)
+    
+    @staticmethod
+    def process_pred_gold(pred, gold):
+        gold = " ".join(gold.split())
+        pred = " ".join(pred.split())
+        return pred, gold
+
+
+class TextDataset(Dataset):
+    def __init__(self, tokenizer, pool, args, file_path, samplenum=-1):
+        self.cnt = 0
+        self.tokenizer = tokenizer
+        self.args = args
+        if isinstance(tokenizer, MyTokenizer):
+            tokenizer_type = "mytok"
+        elif isinstance(tokenizer, T5Tokenizer):
+            tokenizer_type = ""
+        elif isinstance(tokenizer, RobertaTokenizer):
+            tokenizer_type = "rb"
+        else:
+            tokenizer_type = "unk"
+        savep = file_path.replace(".jsonl", tokenizer_type + ".exps")
+        # savep = "/home/v-zhuoli1/lzzz/processed/chunk_25.exps"
+        if os.path.exists(savep):
+            logger.info("Loading examples from {}".format(savep))
+            examples = torch.load(savep)
+        else:
+            logger.info("Reading examples from {}".format(file_path))
+            examples = read_review_examples(file_path, samplenum, tokenizer)
+            logger.info(f"Tokenize examples: {file_path}")
+            examples = pool.map(self.tokenize, \
+                [(example, tokenizer, args) for example in examples])
+            torch.save(examples, savep)
+        logger.info("Convert examples to features...")
+        self.set_start_end_ids(examples)
+        self.featss = pool.map(self.convert_examples_to_features, \
+            [(example, tokenizer, args) for example in examples])
+        self.feats = [feat for feats in self.featss for feat in feats]  # expand the lists
+
+    def __len__(self):
+        return len(self.feats)
+
+    def __getitem__(self, i):
+        return self.feats[i]
+
+    def reset_len(self, data_len):
+        assert len(self.feats) >= data_len
+        self.feats = self.feats[:data_len]
+
+    def set_start_end_ids(self, examples):
+        for example in examples:
+            labels = example.labels
+            start_id = 0
+            end_id = len(labels) - 1
+            for i, label in enumerate(labels):
+                if label != -100:               # find the first label
+                    start_id = i
+                    break
+            for i in range(len(labels) - 1, -1, -1):
+                label = labels[i]
+                if label != -100:
+                    end_id = i
+                    break
+            example.start_id = start_id
+            example.end_id = end_id
+
+    def tokenize(self, item):
+        example, tokenizer, args = item
+        example.input = self.encode_remove(tokenizer, example.input, args)
+        e0id = tokenizer.special_dict["<e0>"]
+        inputs = " ".join(str(id) for id in example.input)
+        lines = inputs.split(" " + str(e0id) + " ")
+        lines = [
+            [int(v) for v in line.split(" ") if len(v) > 0] for line in lines
+        ]
+        lens = [len(line) for line in lines]
+        # if 0 in lens:
+        #     logger.info("Warning: empty line in an example.")
+        lens = list(map(len, lines))
+        curlen = len(lens) + sum(lens)
+        left, right = 0, len(lines)
+        while curlen > args.max_source_length - 2:
+            if left % 2 == 0:
+                curlen -= 1 + len(lines[left])
+                left += 1
+            else:
+                right -= 1
+                curlen -= 1 + len(lines[right])
+        lines = lines[left:right]
+        labels = example.labels[left:right]
+        assert len(lines) + sum(map(len, lines)) <= args.max_source_length - 2, "Too long inputs in TextDataset.tokenize."
+        if len(lines) != len(labels):
+            logger.info("Not equal length in TextDataset.tokenize.")
+            lines = lines[:len(labels)]
+            labels = labels[:len(lines)]
+        example.lines = lines
+        example.labels = labels
+        example.msg = self.encode_remove(tokenizer, example.msg, args)
+        return example
+
+    def convert_examples_to_features(self, item):
+        example, _, _ = item
+        if len(example.msg) > 0:
+            exs = []
+            for _ in range(3):  # up sampling
+                if random.random() < 0.5:
+                    exs.append(self.genmsg_example(item))
+                else:
+                    exs.append(self.daemsg_example(item))
+            return exs
+        if random.random() < 0.5:
+            return [self.encoder_example(item)]
+        return [self.decoder_example(item)]
+
+    def encoder_example(self, item):
+        example, tokenizer, args = item
+        lines = example.lines
+        labels = example.labels
+        target_ids = [tokenizer.pad_id] * args.max_target_length
+        source_ids, input_labels = [], []
+        for i, (line, label) in enumerate(zip(lines, labels)):
+            if i == example.start_id:
+                source_ids.append(tokenizer.start_id)
+                input_labels.append(-100)
+            if label != -100:       # only insert special tokens at diffs, not context
+                source_ids.append(tokenizer.mask_id)
+                input_labels.append(label)
+            source_ids.extend(line)
+            input_labels.extend([-100] * len(line))
+            if i == example.end_id:
+                source_ids.append(tokenizer.end_id)
+                input_labels.append(-100)
+        assert len(input_labels) == len(source_ids), "Not equal length."
+        assert len(input_labels) <= args.max_source_length, f"Too long inputs: {len(input_labels)}."
+        source_ids = source_ids[:args.max_source_length - 2]
+        input_labels = input_labels[:args.max_source_length - 2]
+        source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
+        input_labels = [-100] + input_labels + [-100]
+        pad_len = args.max_source_length - len(source_ids)
+        source_ids += [tokenizer.pad_id] * pad_len
+        input_labels += [-100] * pad_len
+
+        new_input_labels = []
+        map_dict = {0: tokenizer.del_id, 1: tokenizer.add_id, 2: tokenizer.keep_id}
+        for label in input_labels:
+            if label == -100:
+                new_input_labels.append(-100)
+            else:
+                new_input_labels.append(map_dict[label])
+        input_labels = new_input_labels
+        assert len(source_ids) == args.max_source_length, "Not equal length."
+        assert len(input_labels) == args.max_source_length, "Not equal length."
+        return ReviewFeatures(example.idx, source_ids, input_labels, target_ids, type="label")
+
+    def decoder_example(self, item):
+        example, tokenizer, args = item
+        lines = example.lines
+        labels = example.labels
+
+        input_labels = [-100] * args.max_source_length
+        source_ids, target_ids = [], []
+        SPECIAL_ID = 0
+        mask_idxs = random.choices(range(len(lines)), k=int(len(lines) * args.mask_rate))
+        id_dict = {0: tokenizer.del_id, 1: tokenizer.add_id, 2: tokenizer.keep_id}
+        for i, (line, label) in enumerate(zip(lines, labels)):
+            if i == example.start_id:
+                source_ids.append(tokenizer.start_id)
+            if label in id_dict:
+                source_ids.append(id_dict[label])
+            if i in mask_idxs:
+                source_ids.append(tokenizer.special_dict[f"<e{SPECIAL_ID}>"])
+                target_ids.append(tokenizer.special_dict[f"<e{SPECIAL_ID}>"])
+                target_ids.extend(line)
+                if SPECIAL_ID < 99:     # only 0-99 ids in vocab
+                    SPECIAL_ID += 1
+            else:
+                source_ids.extend(line)
+            if i == example.end_id:
+                source_ids.append(tokenizer.end_id)
+        source_ids, target_ids = self.pad_assert(source_ids, target_ids, args, tokenizer)
+        return ReviewFeatures(example.idx, source_ids, input_labels, target_ids, type="line")
+
+    def genmsg_example(self, item):
+        example, tokenizer, args = item
+        lines = example.lines
+        labels = example.labels
+        input_labels = [-100] * args.max_source_length
+        source_ids, target_ids = [], []
+        id_dict = {0: tokenizer.del_id, 1: tokenizer.add_id, 2: tokenizer.keep_id}
+        for i, (line, label) in enumerate(zip(lines, labels)):
+            if i == example.start_id:
+                source_ids.append(tokenizer.start_id)
+            if label != -100:
+                source_ids.append(id_dict[label])
+            source_ids.extend(line)
+            if i == example.end_id:
+                source_ids.append(tokenizer.end_id)
+        target_ids.append(tokenizer.msg_id)
+        target_ids.extend(example.msg)
+        assert len(source_ids) <= args.max_source_length, f"Too long inputs: {len(source_ids)}."
+        source_ids, target_ids = self.pad_assert(source_ids, target_ids, args, tokenizer)
+        return ReviewFeatures(example.idx, source_ids, input_labels, target_ids, type="genmsg")
+
+    def daemsg_example(self, item):
+        example, tokenizer, args = item
+        input_labels = [-100] * args.max_source_length
+        source_ids, target_ids = [], []
+        msg_ids = cp(example.msg)
+        masks = [random.random() < 0.20 for _ in range(len(msg_ids))]
+        if sum(masks) == 0:
+            idx = random.choice(range(len(msg_ids)))
+            masks[idx] = True
+        source_ids, target_ids = [], []
+        i = 0
+        SPECIAL_ID = 0
+        while i < len(masks):
+            j = i
+            while j < len(masks) and not masks[j]:
+                source_ids.append(msg_ids[j])
+                j += 1
+            if j == len(masks):
+                break
+            source_ids.append(tokenizer.special_dict[f"<e{SPECIAL_ID}>"])
+            target_ids.append(tokenizer.special_dict[f"<e{SPECIAL_ID}>"])
+            while j < len(masks) and masks[j]:
+                target_ids.append(msg_ids[j])
+                j += 1
+            if SPECIAL_ID < 99:     # only 0-99 ids in vocab
+                SPECIAL_ID += 1
+            i = j
+        source_ids, target_ids = self.pad_assert(source_ids, target_ids, args, tokenizer)
+        return ReviewFeatures(example.idx, source_ids, input_labels, target_ids, type="daemsg")
+
+    def pad_assert(self, source_ids, target_ids, args, tokenizer):
+        source_ids = source_ids[:args.max_source_length - 2]
+        source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
+        pad_len = args.max_source_length - len(source_ids)
+        source_ids += [tokenizer.pad_id] * pad_len
+        target_ids = target_ids[:args.max_target_length - 1]
+        target_ids = target_ids + [tokenizer.eos_id]
+        pad_len = args.max_target_length - len(target_ids)
+        target_ids += [tokenizer.pad_id] * pad_len
+        assert len(source_ids) == args.max_source_length, "Not equal length."
+        assert len(target_ids) == args.max_target_length, "Not equal length."
+        return source_ids, target_ids
+
+    def encode_remove(self, tokenizer, text, args):
+        text = tokenizer.encode(text, max_length=args.max_source_length, truncation=True)
+        if type(tokenizer) == T5Tokenizer:
+            return text[:-1]
+        elif type(tokenizer) == RobertaTokenizer:
+            return text[1:-1]
+        elif type(tokenizer) == MyTokenizer:
+            return text
+        else:
+            raise NotImplementedError
+
+
+class CommentGenDataset(TextDataset):
+    def __init__(self, tokenizer, pool, args, file_path, samplenum=-1):
+        self.tokenizer = tokenizer
+        if isinstance(tokenizer, MyTokenizer):
+            tokenizer_type = "mytok"
+        elif isinstance(tokenizer, T5Tokenizer):
+            tokenizer_type = ""
+        elif isinstance(tokenizer, RobertaTokenizer):
+            tokenizer_type = "rb"
+        else:
+            tokenizer_type = "unk"
+        savep = file_path.replace(".jsonl", tokenizer_type + ".exps")
+        if os.path.exists(savep):
+            logger.info("Loading examples from {}".format(savep))
+            examples = torch.load(savep)
+        else:
+            logger.info("Reading examples from {}".format(file_path))
+            examples = read_review_examples(file_path, samplenum, tokenizer)
+            # for i in range(len(examples)):
+            #     examples[i].msg = " ".join(nltk.word_tokenize(examples[i].msg))
+            logger.info(f"Tokenize examples: {file_path}")
+            examples = pool.map(self.tokenize, \
+                [(example, tokenizer, args) for example in examples])
+            torch.save(examples, savep)
+        logger.info("Convert examples to features...")
+        self.set_start_end_ids(examples)
+        self.feats = pool.map(self.convert_examples_to_features, \
+            [(example, tokenizer, args) for example in examples])
+        self.feats = [feat for feat in self.feats if feat is not None]
+
+    def convert_examples_to_features(self, item):
+        example, tokenizer, args = item
+        if len(example.msg) == 0:
+            return None
+        return self.genmsg_example(item)
+
+
+class CommentClsDataset(TextDataset):
+    def __init__(self, tokenizer, pool, args, file_path, samplenum=-1):
+        self.tokenizer = tokenizer
+        if isinstance(tokenizer, MyTokenizer):
+            tokenizer_type = "mytok"
+        elif isinstance(tokenizer, T5Tokenizer):
+            tokenizer_type = ""
+        elif isinstance(tokenizer, RobertaTokenizer):
+            tokenizer_type = "rb"
+        else:
+            tokenizer_type = "unk"
+        savep = file_path.replace(".jsonl", tokenizer_type + ".exps")
+        if os.path.exists(savep):
+            logger.info("Loading examples from {}".format(savep))
+            examples = torch.load(savep)
+        else:
+            logger.info("Reading examples from {}".format(file_path))
+            examples = read_review_examples(file_path, samplenum, tokenizer)
+            logger.info(f"Tokenize examples: {file_path}")
+            examples = pool.map(self.tokenize, \
+                [(example, tokenizer, args) for example in examples])
+            torch.save(examples, savep)
+        logger.info("Convert examples to features...")
+        self.set_start_end_ids(examples)
+        self.feats = pool.map(self.convert_examples_to_features, \
+            [(example, tokenizer, args) for example in examples])
+
+    def convert_examples_to_features(self, item):
+        example, tokenizer, args = item
+        tmpfeature = self.genmsg_example(item)
+        return ClsFeatures(tmpfeature.example_id, tmpfeature.source_ids, example.y)
+
+
+class SimpleClsDataset(TextDataset):
+    def __init__(self, tokenizer, pool, args, file_path, samplenum=-1):
+        self.tokenizer = tokenizer
+        if isinstance(tokenizer, MyTokenizer):
+            tokenizer_type = "mytok"
+        elif isinstance(tokenizer, T5Tokenizer):
+            tokenizer_type = ""
+        elif isinstance(tokenizer, RobertaTokenizer):
+            tokenizer_type = "rb"
+        else:
+            tokenizer_type = "unk"
+        savep = file_path.replace(".jsonl", tokenizer_type + ".simpexps")
+        if os.path.exists(savep):
+            logger.info("Loading examples from {}".format(savep))
+            self.feats = torch.load(savep)
+        else:
+            logger.info("Reading examples from {}".format(file_path))
+            examples = read_review_examples(file_path, samplenum, tokenizer)
+            logger.info(f"Tokenize examples: {file_path}")
+            self.feats = pool.map(self.convert_examples_to_features, \
+                [(example, tokenizer, args) for example in examples])
+            torch.save(self.feats, savep)
+
+    def convert_examples_to_features(self, item):
+        example, tokenizer, args = item
+        example.input_lines = example.input.split("<e0>")
+        labels_l = len(example.labels)
+        example.input_lines = example.input_lines[:labels_l]
+        for i in range(len(example.input_lines)):
+            if example.labels[i] == 1:
+                example.input_lines[i] = "+ " + example.input_lines[i]
+            elif example.labels[i] == 0:
+                example.input_lines[i] = "- " + example.input_lines[i]
+        example.input = " ".join(example.input_lines)
+        input_ids = self.encode_remove(tokenizer, example.input, args)
+        exceed_l = len(input_ids) - args.max_source_length + 2
+        if exceed_l > 0:
+            halfexl = (exceed_l + 1) // 2
+            input_ids = input_ids[halfexl:-halfexl]
+        source_ids = input_ids[:args.max_source_length - 2]
+        source_ids = [tokenizer.bos_id] + source_ids + [tokenizer.eos_id]
+        pad_len = args.max_source_length - len(source_ids)
+        source_ids += [tokenizer.pad_id] * pad_len
+        example_id = example.idx
+        y = example.y
+        return ClsFeatures(example_id, source_ids, y)
+
+
+class SimpleGenDataset(TextDataset):
+    def __init__(self, tokenizer, pool, args, file_path, samplenum=-1):
+        self.tokenizer = tokenizer
+        if isinstance(tokenizer, MyTokenizer):
+            tokenizer_type = "mytok"
+        elif isinstance(tokenizer, T5Tokenizer):
+            tokenizer_type = ""
+        elif isinstance(tokenizer, RobertaTokenizer):
+            tokenizer_type = "rb"
+        else:
+            tokenizer_type = "unk"
+        savep = file_path.replace(".jsonl", tokenizer_type + ".simpgenexps")
+        if os.path.exists(savep):
+            logger.info("Loading examples from {}".format(savep))
+            self.feats = torch.load(savep)
+        else:
+            logger.info("Reading examples from {}".format(file_path))
+            data = read_jsonl(file_path)
+            # data = [dic for dic in data if len(dic["patch"].split("\n")) <= 20]
+            for i in range(len(data)):
+                data[i]["idx"] = i
+            logger.info(f"Tokenize examples: {file_path}")
+            # self.feats = pool.map(self.convert_examples_to_features, \
+            #     [(dic, tokenizer, args) for dic in data])
+            self.feats = [self.convert_examples_to_features((dic, tokenizer, args)) for dic in data]
+            torch.save(self.feats, savep)
+
+    def convert_examples_to_features(self, item):
+        dic, tokenizer, args = item
+        diff, msg = dic["patch"], dic["msg"]
+        difflines = diff.split("\n")[1:]        # remove start @@
+        difflines = [line for line in difflines if len(line.strip()) > 0]
+        map_dic = {"-": 0, "+": 1, " ": 2}
+        def f(s):
+            if s in map_dic:
+                return map_dic[s]
+            else:
+                return 2
+        labels = [f(line[0]) for line in difflines]
+        difflines = [line[1:].strip() for line in difflines]
+        inputstr = ""
+        for label, line in zip(labels, difflines):
+            if label == 1:
+                inputstr += "<add>" + line
+            elif label == 0:
+                inputstr += "<del>" + line
+            else:
+                inputstr += "<keep>" + line
+        source_ids = self.encode_remove(tokenizer, inputstr, args)
+        target_ids = []
+        target_ids.append(tokenizer.msg_id)
+        msg = self.encode_remove(tokenizer, dic["msg"], args)
+        target_ids.extend(msg)
+        source_ids, target_ids = self.pad_assert(source_ids, target_ids, args, tokenizer)
+        input_labels = [-100] * len(source_ids)
+        return ReviewFeatures(dic["idx"], source_ids, input_labels, target_ids, type="genmsg")
+
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+
+    def __init__(self, example_id, source_ids, target_ids, url=None):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.target_ids = target_ids
+        self.url = url
+
+
+class ReviewFeatures(object):
+    def __init__(self, example_id, source_ids, source_labels, target_ids, type):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.source_labels = source_labels
+        self.target_ids = target_ids
+        assert type in ("label", "line", "genmsg", "daemsg")
+        self.type = type
+
+class ClsFeatures(object):
+    def __init__(self, example_id, source_ids, y):
+        self.example_id = example_id
+        self.source_ids = source_ids
+        self.y = y
+
+class ReviewExample(object):
+    """A single training/test example."""
+
+    def __init__(
+        self, idx, oldf, diff, msg, cmtid, max_len, y
+    ):
+        self.idx = idx      # idx is useless yet
+        self.oldf = oldf
+        self.diff = diff
+        self.msg = msg
+        self.cmtid = cmtid
+        self.max_len = max_len
+        self.y = y
+        self.prevlines = []
+        self.afterlines = []
+        self.lines = []
+        self.labels = []
+        self.avail = False
+        self.input = ""
+        self.align_and_clean()
+        self.postprocess()
+
+    def postprocess(self):
+        if not self.avail:
+            return
+        # Warning: lines is not self.lines
+        # lines for rough length estimation
+        lines = [source_str.split() for source_str in self.lines]
+        inputl = len(lines) # line tag
+        inputl += sum(map(len, lines))
+        left, right = 0, len(lines)
+        while inputl > self.max_len:
+            if left % 2 == 0:
+                inputl -= len(lines[left]) + 1
+                left += 1
+            else:
+                right -= 1
+                inputl -= len(lines[right]) + 1
+        lines = lines[left:right]
+        self.lines = self.lines[left:right]
+        self.labels = self.labels[left:right]
+        prevlines = self.prevlines
+        afterlines = self.afterlines
+        prev_after_len = max(len(prevlines), len(afterlines))
+        i = 0
+        while inputl < self.max_len and i < prev_after_len:
+            if i < len(prevlines):
+                newl = inputl + len(prevlines[-1-i].split()) + 1
+                if newl > self.max_len:
+                    break
+                self.lines.insert(0, prevlines[-1-i])
+                self.labels.insert(0, -100)
+                inputl = newl  # tag
+            if i < len(afterlines):
+                newl = inputl + len(afterlines[i].split()) + 1
+                if newl > self.max_len:
+                    break
+                self.lines.append(afterlines[i])
+                self.labels.append(-100)
+                inputl = newl    # tag
+            i += 1
+        assert inputl <= self.max_len, "Too long inputs."
+        assert len(self.lines) == len(self.labels), "Not equal length."
+        self.input = "<e0>".join(self.lines)
+        self.prevlines, self.lines, self.afterlines = [], [], []
+
+    def remove_space_clean(self, line):
+        """
+            Remove start and end empty chars.
+        """
+        rep = " \t\r"
+        totallen = len(line)
+        i = 0
+        while i < totallen and line[i] in rep:
+            i += 1
+        j = totallen - 1
+        while j >= 0 and line[j] in rep:
+            j -= 1
+        line = line[i : j + 1]
+        return line
+
+    def align_and_clean(self):
+        oldflines = self.oldf.split("\n")
+        difflines = self.diff.split("\n")
+        first_line = difflines[0]
+        difflines = difflines[1:]
+        difflines = [line for line in difflines if line != r""]
+        regex = r"@@ -(\d+),(\d+) \+(\d+),(\d+) @@"
+        matchres = re.match(regex, first_line)
+        if matchres:
+            startline, rangelen, startpos, endpos = matchres.groups()
+            self.avail = True
+        else:
+            self.avail = False
+            return
+        startline, rangelen = int(startline) - 1, int(rangelen)
+        endline = startline + rangelen
+        self.prevlines = oldflines[:startline]
+        self.afterlines = oldflines[endline:]
+        for line in difflines:
+            if line.startswith("-"):
+                self.lines.append(line[1:])
+                self.labels.append(0)
+            elif line.startswith("+"):
+                self.lines.append(line[1:])
+                self.labels.append(1)
+            else:
+                self.lines.append(line)
+                self.labels.append(2)
+        self.prevlines = [self.remove_space_clean(line) for line in self.prevlines]
+        self.afterlines = [self.remove_space_clean(line) for line in self.afterlines]
+        self.lines = [self.remove_space_clean(line) for line in self.lines]
+        self.msg = self.remove_space_clean(self.msg)
+        self.prevlines = [line for line in self.prevlines if len(line) > 0]
+        self.afterlines = [line for line in self.afterlines if len(line) > 0]
+        # print("\n".join(self.prevlines))
+        # print("\n\n\n\n")
+        # print("\n".join(self.lines))
+        # print("\n\n\n\n")
+        # print("\n".join(self.afterlines))
+        # print("\n\n\n\n")
+        assert len(self.lines) == len(self.labels), "Not equal length in align."
+        topack = list(
+            zip(
+                *[
+                    (line, label)
+                    for line, label in zip(self.lines, self.labels)
+                    if len(line) > 0
+                ]
+            )
+        )
+        if topack == []:
+            self.avail = False
+            return
+        else:
+            self.lines, self.labels = topack
+        # tuple->list, convenient for later operation
+        self.lines = list(self.lines)
+        self.labels = list(self.labels)
+
+
+def read_review_examples(filename, data_num=-1, tokenizer=None):
+    """Read examples from filename."""
+    examples = []
+    idx = 0
+    with open(filename) as f:
+        for line in f:
+            try:
+                js = json.loads(line.strip())
+            except:
+                print("Error during reading json data.")
+                continue
+            maxl = 200
+            if "y" not in js:
+                js["y"] = 0
+            if "msg" in js and len(js["msg"]) > 0:
+                js["y"] = 1
+            example = ReviewExample(
+                        idx=idx,
+                        oldf=js["oldf"],
+                        diff=js["patch"],
+                        msg=js["msg"] if "msg" in js else "",
+                        cmtid=js["cmtid"] if "cmtid" in js else "",
+                        max_len=maxl,
+                        y=js["y"]
+                    )
+            if example.avail:
+                examples.append(example)
+                idx += 1
+                if idx == data_num:
+                    break
+            else:
+                # print(f"Passing {idx} because of invalid diff.")
+                idx += 1 
+                if idx == data_num:
+                    break
+                
+    return examples
+
+
+def read_jsonl(path):
+    data = []
+    with open(path) as f:
+        for line in f:
+            try:
+                js = json.loads(line.strip())
+            except:
+                print("Error during reading json data.")
+                continue
+            data.append(js)
+    return data