python sample generation script

coreml_example.py

@@ -0,0 +1,90 @@
+import time
+import numpy as np
+from argparse import ArgumentParser
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
+
+parser = ArgumentParser()
+parser.add_argument("--model_path", "--model-path", required=True)
+parser.add_argument("--prompt", "-p", required=True)
+parser.add_argument("--max-tokens", "--max_tokens", type=int, default=100)
+parser.add_argument("--min_p", "--min-p", type=float, default=0.3)
+parser.add_argument("--temp", type=float, default=1.0)
+args = parser.parse_args()
+
+import coremltools as ct
+
+print("Loading model...")
+if args.model_path.rstrip("/").endswith(".mlpackage"):
+    mf_model_1 = ct.models.MLModel(
+        args.model_path,
+        compute_units=ct.ComputeUnit.CPU_AND_NE,
+        function_name="length_1",
+    )
+    mf_model_64 = ct.models.MLModel(
+        args.model_path,
+        compute_units=ct.ComputeUnit.CPU_AND_NE,
+        function_name="length_64",
+    )
+else:
+    mf_model_1 = ct.models.CompiledMLModel(
+        args.model_path,
+        compute_units=ct.ComputeUnit.CPU_AND_NE,
+        function_name="length_1",
+    )
+    mf_model_64 = ct.models.CompiledMLModel(
+        args.model_path,
+        compute_units=ct.ComputeUnit.CPU_AND_NE,
+        function_name="length_64",
+    )
+
+
+def min_p_sample(logits, min_p, temp):
+    # logits = logits.astype(np.float16)
+    max_ = np.max(logits * (1 / temp), axis=1, keepdims=True)
+    logits = logits - max_
+    logits = np.exp(logits)
+    logits[logits < min_p] = 0
+    # logits = logits.astype(np.float32)
+    logits = np.cumsum(logits, axis=1)
+    sample = np.random.uniform(high=logits[:, -1:])
+    sample = np.argmax(logits > sample, axis=1).astype(np.int32)
+    return sample
+
+
+length = len(tokenizer(args.prompt)["input_ids"])
+input_ids = tokenizer(
+    args.prompt, return_tensors="np", padding="max_length", max_length=64
+)["input_ids"].astype(np.int32)
+print("Prompt:", args.prompt)
+state = mf_model_64.make_state()
+start = time.time()
+pred = mf_model_64.predict(
+    {"input_ids": input_ids, "query_pos1": np.array([0], dtype=np.int32)}, state
+)
+prompt_time = time.time() - start
+# input_ids = pred["logits"][..., length - 1].argmax(1, keepdims=True).astype(np.int32)
+logits = pred["logits"][..., [length - 1]]
+input_ids = min_p_sample(logits, args.min_p, args.temp)
+print("Generated:")
+print(tokenizer.decode(input_ids[0]), end="", flush=True)
+start = time.time()
+for i in range(args.max_tokens):
+    pred = mf_model_1.predict(
+        {"input_ids": input_ids, "query_pos1": np.array([i + length], dtype=np.int32)},
+        state,
+    )
+    input_ids = min_p_sample(pred["logits"], args.min_p, args.temp)
+    # input_ids = pred["logits"].argmax(1).astype(np.int32)
+    print(tokenizer.decode(input_ids[0]), end="", flush=True)
+print("", "=" * 10)
+generation_time = time.time() - start
+
+print(
+    "Prompt:",
+    length / prompt_time,
+    "tokens-per-sec",
+    f"({64 / prompt_time} considering the processed padding)",
+)
+print("Generation:", args.max_tokens / generation_time, "tokens-per-sec")