<reponame>serjtroshin/PLBART
import argparse
import logging
import os
from pathlib import Path

logging.basicConfig(level=logging.INFO)

from submit import submit_local


def run_generate(args):
    logging.info("run_generate")
    model = args.save_dir + f"/{args.model_name}"
    FILE_PREF = args.save_dir + "/output"
    additional = ""
    args.debug = True
    if args.debug:
        additional += "--max-update 10 --max-epoch 1"
    script = f"""
    fairseq-generate {args.data_dir}/data-bin \
        --path {model} \
        --task translation_from_pretrained_bart \
        --gen-subset test \
        -t {args.target} -s {args.source} \
        --sacrebleu \
        --remove-bpe 'sentencepiece' \
        --batch-size 4 \
        --langs {args.langs} \
        --beam 10 \
        --lenpen 1.0 > {FILE_PREF}
    """.strip()
    submit_local(script)

def run_cat(args):
    logging.info("run_cat")
    FILE_PREF = args.save_dir + "/output"
    script = f"""
    cat {FILE_PREF} | grep -P "^H" |sort -V |cut -f 3- | sed 's/\[{args.target}\]//g' > {FILE_PREF}.hyp
    """.strip()
    submit_local(script)
    
def run_evaluator(args):
    logging.info("run_evaluator")
    GOUND_TRUTH_PATH = args.data_dir + "/test.json"
    FILE_PREF = args.save_dir + "/output"
    RESULT_FILE = args.save_dir + "/result.txt"
    parent_dir = os.path.dirname(os.path.abspath(__file__))

    print("CodeXGlue Evaluation")
    script = f"""
    python3 {parent_dir}/evaluator.py \
        --expected {GOUND_TRUTH_PATH} \
        --predicted {FILE_PREF}.hyp \
        --fmt_output {args.fmt_output} \
        2>&1 | tee -a {RESULT_FILE};
    """.strip()
    submit_local(script)

def run_eval_script(args):
    logging.info("run_eval_script")
    GOUND_TRUTH_PATH = args.data_dir + "/test.json"
    FILE_PREF = args.save_dir + "/output"
    RESULT_FILE = args.save_dir + "/result.txt"
    CB_EVAL_SCRIPT=args.calc_code_bleu
    CB_EVAL_DIR=os.path.dirname(os.path.abspath(args.calc_code_bleu))

    TARGET=args.target
    script = f"""
    cd {CB_EVAL_DIR} && python3 calc_code_bleu.py \
    --hyp {FILE_PREF}.hyp \
    --refs {GOUND_TRUTH_PATH} \
    --json_refs \
    --lang {TARGET} \
    --fmt_output {args.fmt_output} \
    2>&1 | tee -a {RESULT_FILE};
    """.strip()
    submit_local(script)

def main():
    parser = argparse.ArgumentParser(description='Evaluate leaderboard predictions for code completion (line level).')
    parser.add_argument("--save_dir", required=True, type=str)
    parser.add_argument("--data_dir", required=True, type=str)
    parser.add_argument("-s", "--source", required=True, type=str)
    parser.add_argument("-t", "--target", required=True, type=str)
    parser.add_argument("--langs", required=True, type=str)
    parser.add_argument("--calc_code_bleu", required=True, type=str)
    parser.add_argument("--model_name", required=True, type=str)
    parser.add_argument("--fmt_output", required=True, type=str)
    args = parser.parse_args()

    run_generate(args)
    run_cat(args)
    run_evaluator(args)
    run_eval_script(args)


if __name__=="__main__":
    main()


# function generate () {

# cat $FILE_PREF | grep -P "^H" |sort -V |cut -f 3- | sed 's/\[${TARGET}\]//g' > $FILE_PREF.hyp;

# echo "CodeXGlue Evaluation" > ${RESULT_FILE}
# python evaluator.py \
#     --expected $GOUND_TRUTH_PATH \
#     --predicted $FILE_PREF.hyp \
#     2>&1 | tee -a ${RESULT_FILE};

# echo "CodeBLEU Evaluation" >> ${RESULT_FILE}
# export PYTHONPATH=${HOME_DIR};
# python $CB_EVAL_SCRIPT \
#     --hyp $FILE_PREF.hyp \
#     --refs $GOUND_TRUTH_PATH \
#     --json_refs \
#     --lang $TARGET \
#     2>&1 | tee -a ${RESULT_FILE};

# }
import argparse
import logging
import os
from pathlib import Path

logging.basicConfig(level=logging.INFO)

from submit import submit_local


def run_generate(args):
    logging.info("run_generate")
    model = args.save_dir + f"/{args.model_name}"
    FILE_PREF = args.save_dir + "/output"
    additional = ""
    args.debug = True
    if args.debug:
        additional += "--max-update 10 --max-epoch 1"
    script = f"""
    fairseq-generate {args.data_dir}/data-bin \
        --path {model} \
        --task translation_from_pretrained_bart \
        --gen-subset test \
        -t {args.target} -s {args.source} \
        --sacrebleu \
        --remove-bpe 'sentencepiece' \
        --batch-size 4 \
        --langs {args.langs} \
        --beam 10 \
        --lenpen 1.0 > {FILE_PREF}
    """.strip()
    submit_local(script)

def run_cat(args):
    logging.info("run_cat")
    FILE_PREF = args.save_dir + "/output"
    script = f"""
    cat {FILE_PREF} | grep -P "^H" |sort -V |cut -f 3- | sed 's/\[{args.target}\]//g' > {FILE_PREF}.hyp
    """.strip()
    submit_local(script)
    
def run_evaluator(args):
    logging.info("run_evaluator")
    GOUND_TRUTH_PATH = args.data_dir + "/test.json"
    FILE_PREF = args.save_dir + "/output"
    RESULT_FILE = args.save_dir + "/result.txt"
    parent_dir = os.path.dirname(os.path.abspath(__file__))

    print("CodeXGlue Evaluation")
    script = f"""
    python3 {parent_dir}/evaluator.py \
        --expected {GOUND_TRUTH_PATH} \
        --predicted {FILE_PREF}.hyp \
        --fmt_output {args.fmt_output} \
        2>&1 | tee -a {RESULT_FILE};
    """.strip()
    submit_local(script)

def run_eval_script(args):
    logging.info("run_eval_script")
    GOUND_TRUTH_PATH = args.data_dir + "/test.json"
    FILE_PREF = args.save_dir + "/output"
    RESULT_FILE = args.save_dir + "/result.txt"
    CB_EVAL_SCRIPT=args.calc_code_bleu
    CB_EVAL_DIR=os.path.dirname(os.path.abspath(args.calc_code_bleu))

    TARGET=args.target
    script = f"""
    cd {CB_EVAL_DIR} && python3 calc_code_bleu.py \
    --hyp {FILE_PREF}.hyp \
    --refs {GOUND_TRUTH_PATH} \
    --json_refs \
    --lang {TARGET} \
    --fmt_output {args.fmt_output} \
    2>&1 | tee -a {RESULT_FILE};
    """.strip()
    submit_local(script)

def main():
    parser = argparse.ArgumentParser(description='Evaluate leaderboard predictions for code completion (line level).')
    parser.add_argument("--save_dir", required=True, type=str)
    parser.add_argument("--data_dir", required=True, type=str)
    parser.add_argument("-s", "--source", required=True, type=str)
    parser.add_argument("-t", "--target", required=True, type=str)
    parser.add_argument("--langs", required=True, type=str)
    parser.add_argument("--calc_code_bleu", required=True, type=str)
    parser.add_argument("--model_name", required=True, type=str)
    parser.add_argument("--fmt_output", required=True, type=str)
    args = parser.parse_args()

    run_generate(args)
    run_cat(args)
    run_evaluator(args)
    run_eval_script(args)


if __name__=="__main__":
    main()


# function generate () {

# cat $FILE_PREF | grep -P "^H" |sort -V |cut -f 3- | sed 's/\[${TARGET}\]//g' > $FILE_PREF.hyp;

# echo "CodeXGlue Evaluation" > ${RESULT_FILE}
# python evaluator.py \
#     --expected $GOUND_TRUTH_PATH \
#     --predicted $FILE_PREF.hyp \
#     2>&1 | tee -a ${RESULT_FILE};

# echo "CodeBLEU Evaluation" >> ${RESULT_FILE}
# export PYTHONPATH=${HOME_DIR};
# python $CB_EVAL_SCRIPT \
#     --hyp $FILE_PREF.hyp \
#     --refs $GOUND_TRUTH_PATH \
#     --json_refs \
#     --lang $TARGET \
#     2>&1 | tee -a ${RESULT_FILE};

# }
fairseq-generate {args.data_dir}/data-bin \         --path {model} \         --task translation_from_pretrained_bart \         --gen-subset test \         -t {args.target} -s {args.source} \         --sacrebleu \         --remove-bpe 'sentencepiece' \         --batch-size 4 \         --langs {args.langs} \         --beam 10 \         --lenpen 1.0 > {FILE_PREF} cat {FILE_PREF} | grep -P "^H" |sort -V |cut -f 3- | sed 's/\[{args.target}\]//g' > {FILE_PREF}.hyp python3 {parent_dir}/evaluator.py \         --expected {GOUND_TRUTH_PATH} \         --predicted {FILE_PREF}.hyp \         --fmt_output {args.fmt_output} \         2>&1 | tee -a {RESULT_FILE}; cd {CB_EVAL_DIR} && python3 calc_code_bleu.py \     --hyp {FILE_PREF}.hyp \     --refs {GOUND_TRUTH_PATH} \     --json_refs \     --lang {TARGET} \     --fmt_output {args.fmt_output} \     2>&1 | tee -a {RESULT_FILE}; # function generate () { # cat $FILE_PREF | grep -P "^H" |sort -V |cut -f 3- | sed 's/\[${TARGET}\]//g' > $FILE_PREF.hyp; # echo "CodeXGlue Evaluation" > ${RESULT_FILE} # python evaluator.py \ #     --expected $GOUND_TRUTH_PATH \ #     --predicted $FILE_PREF.hyp \ #     2>&1 | tee -a ${RESULT_FILE}; # echo "CodeBLEU Evaluation" >> ${RESULT_FILE} # export PYTHONPATH=${HOME_DIR}; # python $CB_EVAL_SCRIPT \ #     --hyp $FILE_PREF.hyp \ #     --refs $GOUND_TRUTH_PATH \ #     --json_refs \ #     --lang $TARGET \ #     2>&1 | tee -a ${RESULT_FILE}; # }