import sys
# sys.path.append('/home/mila/x/xiyuan.zou/research/icl-reverse-cot')
# import hydra
# from omegaconf import DictConfig
# from utils import load_dataset, chunks, process_answer
import torch
import transformers
from transformers import AutoTokenizer, AutoModel, GPTJForCausalLM, AutoModelForCausalLM
# import hydra
# from omegaconf import DictConfig, OmegaConf
import argparse
import os

def main(input_file, output_file):
    outputs = []
    with open(input_file, 'r') as reader:
        input_lines = reader.readlines()
    for i, each_line in enumerate(input_lines):
        if i > 1 and input_lines[i - 1].startswith("[0."):
            outputs.append(each_line)

    with open(output_file, 'w') as writer:
        for each in outputs:
            writer.write(each)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_file', dest='input_file', action='store', required=True, help='num training examples to use', type=str)
    parser.add_argument('--output_file', dest='output_file', action='store', required=True, help='num training examples to use', type=str)
    # parser.add_argument('--dataset', dest='dataset', action='store', required=True, help='num training examples to use', type=str)
    # parser.add_argument('--model', dest='model', action='store', required=True, help='num training examples to use', type=str)
    # parser.add_argument('--model', dest='model', action='store', required=True, help='num training examples to use', type=str)
    # compression_token_initialization

    args = parser.parse_args()
    args = vars(args)

    # simple processing
    # def convert_to_list(items, is_int=False):
    #     if is_int:
    #         return [int(s.strip()) for s in items.split(",")]
    #     else:
    #         return [s.strip() for s in items.split(",")]

    # args['models'] = convert_to_list(args['models'])
    # args['datasets'] = convert_to_list(args['datasets'])
    # args['all_shots'] = convert_to_list(args['all_shots'], is_int=True)

    main(**args)