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))
|