codesnippets / get_sample_code.py
patrickvonplaten's picture
up
685ce0f
#!/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