diff --git a/src/transformers/benchmark/cache_benchmark.py b/src/transformers/benchmark/cache_benchmark.py index c59286a745..e60c14ebcd 100644 --- a/src/transformers/benchmark/cache_benchmark.py +++ b/src/transformers/benchmark/cache_benchmark.py @@ -165,7 +165,11 @@ class CacheSpeedBenchMark(SpeedBenchMark, CacheBenchMark): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( - "--config_path", default=None, type=str, required=False, help="Path to a prepared run file or a previously run output file." + "--config_path", + default=None, + type=str, + required=False, + help="Path to a prepared run file or a previously run output file.", ) parser.add_argument( "--output_path", type=str, required=True, help="Path to the output file where the run's info. will be saved." @@ -173,7 +177,6 @@ if __name__ == "__main__": args = parser.parse_args() if args.config_path is None: - repo_id = "meta-llama/Llama-2-7b-hf" prefill_num_iter = 3 num_iter = 16 diff --git a/src/transformers/benchmark/from_pretrained_benchmark.py b/src/transformers/benchmark/from_pretrained_benchmark.py index 2848f456b8..d1a4e15fe8 100644 --- a/src/transformers/benchmark/from_pretrained_benchmark.py +++ b/src/transformers/benchmark/from_pretrained_benchmark.py @@ -38,11 +38,13 @@ class FromPretrainedSpeedBenchMark(SpeedBenchMark, FromPretrainedBenchMark): if __name__ == "__main__": - from transformers import AutoModel - parser = argparse.ArgumentParser() parser.add_argument( - "--config_path", default=None, type=str, required=False, help="Path to a prepared run file or a previously run output file." + "--config_path", + default=None, + type=str, + required=False, + help="Path to a prepared run file or a previously run output file.", ) parser.add_argument( "--output_path", type=str, required=True, help="Path to the output file where the run's info. will be saved."