Add missing import (#497)

* Add missing import

* format

---------

Co-authored-by: Awni Hannun <awni@apple.com>

llms/mlx_lm/generate.py

@@ -115,7 +115,14 @@ def main(args):
     formatter = colorprint_by_t0 if args.colorize else None
 
     generate(
-        model, tokenizer, prompt, args.temp, args.max_tokens, True, formatter=formatter, top_p=args.top_p
+        model,
+        tokenizer,
+        prompt,
+        args.temp,
+        args.max_tokens,
+        True,
+        formatter=formatter,
+        top_p=args.top_p,
     )
 
 

llms/mlx_lm/merge.py

@@ -3,6 +3,7 @@
 import argparse
 import glob
 import json
+import shutil
 from pathlib import Path
 
 import mlx.core as mx