Format bench-TriLMs.py with black
Browse files- bench-TriLMs.py +10 -6
bench-TriLMs.py
CHANGED
@@ -117,7 +117,7 @@ def llama_bench(
|
|
117 |
logger.debug(result.stderr.decode(errors="ignore"))
|
118 |
if result.returncode != 0 or len(result.stdout) == 0:
|
119 |
logger.error("Failed to run %s", " ".join(command))
|
120 |
-
break
|
121 |
|
122 |
new_output = json.loads(result.stdout)
|
123 |
logger.info(json.dumps(new_output, indent=4))
|
@@ -171,7 +171,9 @@ def parse_args(args: Sequence[str]):
|
|
171 |
default=Path(os.path.curdir) / "result.json",
|
172 |
help="Path of the benchmark results to be written",
|
173 |
)
|
174 |
-
parser.add_argument(
|
|
|
|
|
175 |
return parser.parse_args(args[1:])
|
176 |
|
177 |
|
@@ -216,13 +218,15 @@ if __name__ == "__main__":
|
|
216 |
|
217 |
if shutil.which("lscpu") is not None:
|
218 |
logger.info("Getting CPU info")
|
219 |
-
final_result["cpuinfo"] = subprocess.run(
|
220 |
-
|
221 |
-
)
|
222 |
|
223 |
if args.gpu and shutil.which("nvidia-smi") is not None:
|
224 |
logger.info("Getting NVIDIA GPU info")
|
225 |
-
final_result["gpuinfo"] = subprocess.run(
|
|
|
|
|
226 |
|
227 |
logger.info("Writing output to: %s", output_file)
|
228 |
logger.debug("Final results: %s", json.dumps(final_result, indent=4))
|
|
|
117 |
logger.debug(result.stderr.decode(errors="ignore"))
|
118 |
if result.returncode != 0 or len(result.stdout) == 0:
|
119 |
logger.error("Failed to run %s", " ".join(command))
|
120 |
+
break
|
121 |
|
122 |
new_output = json.loads(result.stdout)
|
123 |
logger.info(json.dumps(new_output, indent=4))
|
|
|
171 |
default=Path(os.path.curdir) / "result.json",
|
172 |
help="Path of the benchmark results to be written",
|
173 |
)
|
174 |
+
parser.add_argument(
|
175 |
+
"--force", action="store_true", help="Overwrite the result file without asking"
|
176 |
+
)
|
177 |
return parser.parse_args(args[1:])
|
178 |
|
179 |
|
|
|
218 |
|
219 |
if shutil.which("lscpu") is not None:
|
220 |
logger.info("Getting CPU info")
|
221 |
+
final_result["cpuinfo"] = subprocess.run(
|
222 |
+
["lscpu"], capture_output=True
|
223 |
+
).stdout.decode(encoding="utf-8")
|
224 |
|
225 |
if args.gpu and shutil.which("nvidia-smi") is not None:
|
226 |
logger.info("Getting NVIDIA GPU info")
|
227 |
+
final_result["gpuinfo"] = subprocess.run(
|
228 |
+
["nvidia-smi", "-q"], capture_output=True
|
229 |
+
).stdout.decode(encoding="utf-8")
|
230 |
|
231 |
logger.info("Writing output to: %s", output_file)
|
232 |
logger.debug("Final results: %s", json.dumps(final_result, indent=4))
|