zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-01 12:49:50 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Docs on prompt scaling (#963)

Browse Source 
* docs on prompt scaling

* remove unused var

* nits
This commit is contained in:
Awni Hannun
2024-08-29 15:05:17 -07:00
committed by GitHub
parent 1003a8b2dd
commit b1186e2a81
4 changed files with 49 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
llms/mlx_lm/cache_prompt.py
View File

@@ -56,7 +56,7 @@ def setup_arg_parser():
parser.add_argument(
"--max-kv-size",
type=int,
default=1024,
default=None,
help="Set the maximum key-value cache size",
)
parser.add_argument(
@@ -147,3 +147,7 @@ def main():
metadata["tokenizer_config"] = json.dumps(tokenizer_config)
metadata["max_kv_size"] = str(args.max_kv_size)
mx.save_safetensors(args.kv_cache_file, cache_dict, metadata)
if __name__ == "__main__":
main()