Repository navigation
Expand file tree
/
Copy pathuse_existing_torch.py
More file actions
56 lines (49 loc) · 1.53 KB
/
Copy pathuse_existing_torch.py
File metadata and controls
56 lines (49 loc) · 1.53 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import glob
import sys
# Only strip targeted libraries when checking prefix
TORCH_LIB_PREFIXES = (
# requirements/*.txt/in
"torch=",
"torchvision=",
"torchaudio=",
"torchcodec=",
# pyproject.toml
'"torch =',
'"torchvision =',
'"torchaudio =',
'"torchcodec =',
)
def main(argv):
parser = argparse.ArgumentParser(
description="Strip torch lib requirements to use installed version."
)
parser.add_argument(
"--prefix",
action="store_true",
help="Strip prefix matches only (default: False)",
)
args = parser.parse_args(argv)
for file in (
*glob.glob("requirements/**/*.txt", recursive=True),
*glob.glob("requirements/**/*.in", recursive=True),
*glob.glob("pyproject.toml"),
):
with open(file) as f:
lines = f.readlines()
if "torch" in "".join(lines).lower():
with open(file, "w") as f:
for line in lines:
if (
args.prefix
and not line.lower().strip().startswith(TORCH_LIB_PREFIXES)
or not args.prefix
and "torch" not in line.lower()
):
f.write(line)
else:
print(f">>> removed from {file}:", line.strip())
if __name__ == "__main__":
main(sys.argv[1:])