#!/usr/bin/env python3
import sys
name = sys.argv[1]
processor_class = sys.argv[2]
model_class = sys.argv[3]
checkpoint = sys.argv[4]
mask = sys.argv[5]

with open("/home/patrick/transformers/src/transformers/file_utils.py", "r") as f:
    lines = f.readlines()

format_dict = {"processor_class": processor_class, "model_class": model_class, "checkpoint": checkpoint, "mask": mask}
with open("./example.py", "w") as f:
    f.write("#!/usr/bin/env python3\n")

    is_in = False
    is_in_code = False
    for line in lines:
        if line.strip() == (name + ' = r"""'):
            is_in = True
        if is_in and "```python" in line:
            is_in_code = True
        if is_in_code:
            if ">>>" in line:
                f.write(line.split(">>> ")[-1].format(**format_dict))
            elif "..." in line:
                f.write(line.split("... ")[-1].format(**format_dict))
        if is_in_code and (line.strip() == '"""'):
            is_in = is_in_code = False