Merge branch 'main' into feat/batch_generate

.gitignore

@@ -6,6 +6,9 @@ __pycache__/
 # C extensions
 *.so
 
+# Vim
+*.swp
+
 # Distribution / packaging
 .Python
 build/

llms/README.md

@@ -20,6 +20,31 @@ The `mlx-lm` package also has:
 - [Merging models](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/MERGE.md)
 - [HTTP model serving](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/SERVER.md)
 
+### Quick Start
+
+To generate text with an LLM use:
+
+```bash
+mlx_lm.generate --prompt "Hi!"
+```
+
+To chat with an LLM use:
+
+```bash
+mlx_lm.chat
+```
+
+This will give you a chat REPL that you can use to interact with the LLM. The
+chat context is preserved during the lifetime of the REPL.
+
+Commands in `mlx-lm` typically take command line options which let you specify
+the model, sampling parameters, and more. Use `-h` to see a list of available
+options for a command, e.g.:
+
+```bash
+mlx_lm.generate -h
+```
+
 ### Python API
 
 You can use `mlx-lm` as a module:
@@ -138,7 +163,7 @@ mlx_lm.convert \
 
 ### Long Prompts and Generations 
 
-MLX LM has some tools to scale efficiently to long prompts and generations:
+`mlx-lm` has some tools to scale efficiently to long prompts and generations:
 
 - A rotating fixed-size key-value cache.
 - Prompt caching
@@ -155,24 +180,30 @@ different queries. To cache a prompt use `mlx_lm.cache_prompt`. For example:
 cat prompt.txt | mlx_lm.cache_prompt \
  --model mistralai/Mistral-7B-Instruct-v0.3 \
  --prompt - \
- --kv-cache-file mistral_prompt.safetensors
+ --prompt-cache-file mistral_prompt.safetensors
 ``` 
 
 Then use the cached prompt with `mlx_lm.generate`:
 
 ```
 mlx_lm.generate \
- --kv-cache-file mistral_prompt.safetensors \
+ --prompt-cache-file mistral_prompt.safetensors \
  --prompt "\nSummarize the above text."
 ```
 
 The cached prompt is treated as a prefix to the supplied prompt. Also notice
 when using a cached prompt, the model to use is read from the cache and need
 not be supplied explicitly.
 
+Prompt caching can also be used in the Python API in order to to avoid
+recomputing the prompt. This is useful in multi-turn dialogues or across
+requests that use the same context. See the
+[example](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/examples/chat.py)
+for more usage details.
+
 ### Supported Models
 
-MLX LM supports thousands of Hugging Face format LLMs. If the model you want to
+`mlx-lm` supports thousands of Hugging Face format LLMs. If the model you want to
 run is not supported, file an
 [issue](https://github.com/ml-explore/mlx-examples/issues/new) or better yet,
 submit a pull request.

llms/mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.
 
-__version__ = "0.18.2"
+__version__ = "0.19.1"

llms/mlx_lm/cache_prompt.py

@@ -7,13 +7,14 @@
 
 import mlx.core as mx
 
-from .utils import load, make_kv_caches
+from .models.cache import make_prompt_cache, save_prompt_cache
+from .utils import load
 
 
 def setup_arg_parser():
  """Set up and return the argument parser."""
  parser = argparse.ArgumentParser(
- description="Cache the KV cache of a prompt to be reused with mlx_lm.generate"
+ description="Cache the state of a prompt to be reused with mlx_lm.generate"
  )
  parser.add_argument(
  "--model",
@@ -60,7 +61,9 @@ def setup_arg_parser():
  help="Set the maximum key-value cache size",
  )
  parser.add_argument(
- "--kv-cache-file", help="The file to save the KV caches in", required=True
+ "--prompt-cache-file",
+ help="The file to save the prompt cache in",
+ required=True,
  )
  parser.add_argument(
  "--prompt",
@@ -115,7 +118,7 @@ def main():
  else:
  prompt = args.prompt
 
- cache = make_kv_caches(model, args.max_kv_size)
+ cache = make_prompt_cache(model, args.max_kv_size)
  y = mx.array(tokenizer.encode(prompt))
 
  # Process the prompt
@@ -137,16 +140,12 @@ def main():
  print(f"Peak memory: {mx.metal.get_peak_memory() / 2**30:.3f} GB")
 
  print("Saving...")
- cache_dict = {}
- for i, c in enumerate(cache):
- cache_dict[f"{i}_keys"] = c.state[0][..., : c.offset, :]
- cache_dict[f"{i}_values"] = c.state[1][..., : c.offset, :]
  metadata = {}
  metadata["model"] = args.model
  metadata["chat_template"] = tokenizer.chat_template
  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)
+ print(f"Peak memory: {mx.metal.get_peak_memory() / 2**30:.3f} GB")
+ save_prompt_cache(args.prompt_cache_file, cache, metadata)
 
 
 if __name__ == "__main__":

llms/mlx_lm/chat.py

@@ -0,0 +1,82 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+import json
+
+import mlx.core as mx
+
+from .models.cache import load_prompt_cache, make_prompt_cache, save_prompt_cache
+from .utils import load, stream_generate
+
+DEFAULT_TEMP = 0.0
+DEFAULT_TOP_P = 1.0
+DEFAULT_SEED = 0
+DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+
+
+def setup_arg_parser():
+ """Set up and return the argument parser."""
+ parser = argparse.ArgumentParser(description="Chat with an LLM")
+ parser.add_argument(
+ "--model",
+ type=str,
+ help="The path to the local model directory or Hugging Face repo.",
+ default=DEFAULT_MODEL,
+ )
+ parser.add_argument(
+ "--adapter-path",
+ type=str,
+ help="Optional path for the trained adapter weights and config.",
+ )
+ parser.add_argument(
+ "--temp", type=float, default=DEFAULT_TEMP, help="Sampling temperature"
+ )
+ parser.add_argument(
+ "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
+ )
+ parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
+ parser.add_argument(
+ "--max-kv-size",
+ type=int,
+ help="Set the maximum key-value cache size",
+ default=None,
+ )
+ return parser
+
+
+def main():
+ parser = setup_arg_parser()
+ args = parser.parse_args()
+
+ mx.random.seed(args.seed)
+
+ model, tokenizer = load(
+ args.model,
+ adapter_path=args.adapter_path,
+ tokenizer_config={"trust_remote_code": True},
+ )
+
+ print(f"[INFO] Starting chat sessiong with {args.model}. To exit, enter 'q'.")
+ prompt_cache = make_prompt_cache(model, args.max_kv_size)
+ while True:
+ query = input(">> ")
+ if query == "q":
+ break
+ messages = [{"role": "user", "content": query}]
+ prompt = tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+ )
+ for response in stream_generate(
+ model,
+ tokenizer,
+ prompt,
+ temp=args.temp,
+ top_p=args.top_p,
+ prompt_cache=prompt_cache,
+ ):
+ print(response, flush=True, end="")
+ print()
+
+
+if __name__ == "__main__":
+ main()

llms/mlx_lm/examples/chat.py

@@ -0,0 +1,53 @@
+# Copyright © 2024 Apple Inc.
+
+"""
+An example of a multi-turn chat with prompt caching.
+"""
+
+from mlx_lm import generate, load
+from mlx_lm.models.cache import load_prompt_cache, make_prompt_cache, save_prompt_cache
+
+model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
+
+# Make the initial prompt cache for the model
+prompt_cache = make_prompt_cache(model)
+
+# User turn
+prompt = "Hi my name is <Name>."
+messages = [{"role": "user", "content": prompt}]
+prompt = tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+)
+
+# Assistant response
+response = generate(
+ model,
+ tokenizer,
+ prompt=prompt,
+ verbose=True,
+ temp=0.0,
+ prompt_cache=prompt_cache,
+)
+
+# User turn
+prompt = "What's my name?"
+messages = [{"role": "user", "content": prompt}]
+prompt = tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+)
+
+# Assistant response
+response = generate(
+ model,
+ tokenizer,
+ prompt=prompt,
+ verbose=True,
+ temp=0.0,
+ prompt_cache=prompt_cache,
+)
+
+# Save the prompt cache to disk to reuse it at a later time
+save_prompt_cache("mistral_prompt.safetensors", prompt_cache)
+
+# Load the prompt cache from disk
+prompt_cache = load_prompt_cache("mistral_prompt.safetensors")

llms/mlx_lm/examples/generate_response.py

@@ -1,3 +1,5 @@
+# Copyright © 2024 Apple Inc.
+
 from mlx_lm import generate, load
 
 # Specify the checkpoint