text-generation-inference

1
import subprocess
2
import argparse
3

4

5
def main():
6
    parser = argparse.ArgumentParser()
7
    parser.add_argument("--check", action="store_true")
8

9
    args = parser.parse_args()
10

11
    output = subprocess.check_output(["text-generation-launcher", "--help"]).decode(
12
        "utf-8"
13
    )
14

15
    wrap_code_blocks_flag = "<!-- WRAP CODE BLOCKS -->"
16
    final_doc = f"# Text-generation-launcher arguments\n\n{wrap_code_blocks_flag}\n\n"
17

18
    lines = output.split("\n")
19

20
    header = ""
21
    block = []
22
    for line in lines:
23
        if line.startswith("  -") or line.startswith("      -"):
24
            rendered_block = "\n".join(block)
25
            if header:
26
                final_doc += f"## {header}\n```shell\n{rendered_block}\n```\n"
27
            else:
28
                final_doc += f"```shell\n{rendered_block}\n```\n"
29
            block = []
30
            tokens = line.split("<")
31
            if len(tokens) > 1:
32
                header = tokens[-1][:-1]
33
            else:
34
                header = line.split("--")[-1]
35
            header = header.upper().replace("-", "_")
36

37
        block.append(line)
38

39
    rendered_block = "\n".join(block)
40
    final_doc += f"## {header}\n```shell\n{rendered_block}\n```\n"
41
    block = []
42

43
    filename = "docs/source/basic_tutorials/launcher.md"
44
    if args.check:
45
        with open(filename, "r") as f:
46
            doc = f.read()
47
            if doc != final_doc:
48
                tmp = "launcher.md"
49
                with open(tmp, "w") as g:
50
                    g.write(final_doc)
51
                diff = subprocess.run(
52
                    ["diff", tmp, filename], capture_output=True
53
                ).stdout.decode("utf-8")
54
                print(diff)
55
                raise Exception(
56
                    "Doc is not up-to-date, run `python update_doc.py` in order to update it"
57
                )
58
    else:
59
        with open(filename, "w") as f:
60
            f.write(final_doc)
61

62

63
if __name__ == "__main__":
64
    main()
65