Awni Hannun
514 lines
19 KiB
Python
514 lines
19 KiB
Python
# Copyright © 2024 Apple Inc.
|
|
|
|
"""
|
|
Adapted from a PyTorch implementation by David Grangier
|
|
"""
|
|
|
|
import argparse
|
|
import collections
|
|
import copy
|
|
import json
|
|
import logging
|
|
import os
|
|
from importlib.metadata import version
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Optional
|
|
|
|
import lm_eval
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
import numpy as np
|
|
from lm_eval.api.model import LM
|
|
from lm_eval.api.registry import register_model
|
|
from tqdm import tqdm
|
|
|
|
from .generate import batch_generate
|
|
from .models.cache import make_prompt_cache
|
|
from .sample_utils import make_sampler
|
|
from .utils import load
|
|
|
|
DEFAULT_MAX_TOKENS = 8192
|
|
|
|
|
|
def _rstrip_until(s, untils):
|
|
"""Limit a string <s> to the first occurrence of any substring in untils."""
|
|
l = len(s)
|
|
f = [s.find(u) for u in untils]
|
|
f = [l if x < 0 else x for x in f]
|
|
return s[: min(f)]
|
|
|
|
|
|
def _lstrip(s, pattern):
|
|
"""Truncate the prefix of the string after the first occurrence of pattern."""
|
|
if (idx := s.find(pattern)) != -1:
|
|
return s[idx + len(pattern) :]
|
|
return s
|
|
|
|
|
|
def _pad_inputs(inputs):
|
|
lengths = np.array([len(x) for x in inputs])
|
|
maxlen = lengths.max()
|
|
padded = np.stack(
|
|
[np.pad(x, (0, maxlen - len(x))) for x in inputs],
|
|
axis=0,
|
|
)
|
|
return mx.array(padded), mx.array(lengths)
|
|
|
|
|
|
def chat_template_fn(**extra_kwargs):
|
|
def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
|
|
return self.tokenizer.apply_chat_template(
|
|
chat_history,
|
|
tokenize=False,
|
|
add_generation_prompt=add_generation_prompt,
|
|
continue_final_message=not add_generation_prompt,
|
|
**extra_kwargs,
|
|
)
|
|
|
|
return apply_chat_template
|
|
|
|
|
|
@register_model("mlxlm")
|
|
class MLXLM(LM):
|
|
|
|
apply_chat_template = chat_template_fn()
|
|
|
|
def __init__(
|
|
self,
|
|
path_or_hf_repo: str,
|
|
max_tokens: Optional[int] = None,
|
|
batch_size: int = 8,
|
|
use_chat_template: Optional[bool] = None,
|
|
trust_remote_code: bool = False,
|
|
sampler: Optional[Callable[[mx.array], mx.array]] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
tokenizer_config = {"trust_remote_code": True if trust_remote_code else None}
|
|
self._model, self.tokenizer = load(
|
|
path_or_hf_repo, tokenizer_config=tokenizer_config
|
|
)
|
|
self._max_tokens = max_tokens
|
|
self._batch_size = batch_size
|
|
self.use_chat_template = use_chat_template
|
|
if use_chat_template is None:
|
|
self.use_chat_template = self.tokenizer.chat_template is not None
|
|
self._sampler = sampler
|
|
|
|
def _process_prompt(self, prompt, step_size: int = 2048):
|
|
prompt = mx.array(prompt)[None]
|
|
cache = make_prompt_cache(self._model)
|
|
for i in range(0, prompt.shape[1], step_size):
|
|
logits = self._model(prompt[:, i : i + step_size], cache=cache)
|
|
mx.eval([c.state for c in cache])
|
|
mx.clear_cache()
|
|
logprobs = nn.log_softmax(logits[:, -1, :].astype(mx.float32))
|
|
return logprobs, cache
|
|
|
|
def _score_fn(self, inputs, cache: Optional[Any] = None, step_size: int = 2048):
|
|
inputs, lengths = _pad_inputs(inputs)
|
|
inputs, targets = inputs[..., :-1], inputs[..., 1:]
|
|
|
|
cache = cache or make_prompt_cache(self._model)
|
|
offset = 0
|
|
scores, is_greedy = [], []
|
|
for i in range(0, inputs.shape[1], step_size):
|
|
inp = inputs[:, i : i + step_size]
|
|
T = inp.shape[1]
|
|
|
|
logits = self._model(inp, cache=cache)
|
|
log_probs = nn.log_softmax(logits.astype(mx.float32))
|
|
|
|
score = mx.take_along_axis(
|
|
log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
|
|
)[..., 0]
|
|
|
|
ig = targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
|
|
ig = mx.where(mx.arange(offset, T + offset) < lengths[:, None], ig, False)
|
|
|
|
mx.eval(score, ig)
|
|
mx.clear_cache()
|
|
|
|
is_greedy.append(ig)
|
|
scores.append(score)
|
|
offset += T
|
|
|
|
scores = mx.concatenate(scores, axis=1)
|
|
is_greedy = mx.concatenate(is_greedy, axis=1)
|
|
|
|
return scores, lengths, is_greedy
|
|
|
|
def _tokenize(self, texts):
|
|
return [
|
|
tuple(
|
|
self.tokenizer.encode(t, add_special_tokens=not self.use_chat_template)
|
|
)
|
|
for t in texts
|
|
]
|
|
|
|
@property
|
|
def tokenizer_name(self) -> str:
|
|
return self.tokenizer.name_or_path.replace("/", "__")
|
|
|
|
def loglikelihood(self, requests) -> list[tuple[float, bool]]:
|
|
"""Compute log-likelihood of generating a continuation from a context.
|
|
Downstream tasks should attempt to use loglikelihood instead of other
|
|
LM calls whenever possible.
|
|
:param requests: list[Instance]
|
|
A list of Instance objects, with property `args` which returns a tuple (context, continuation).
|
|
`context: str`
|
|
Context string. Implementations of LM must be able to handle an
|
|
empty context string.
|
|
`continuation: str`
|
|
The continuation over which log likelihood will be calculated. If
|
|
there is a word boundary, the space should be in the continuation.
|
|
For example, context="hello" continuation=" world" is correct.
|
|
:return: list[tuple[float, bool]]
|
|
A list of pairs (logprob, isgreedy)
|
|
`logprob: float`
|
|
The log probability of `continuation`.
|
|
`isgreedy`:
|
|
Whether `continuation` would be generated by greedy sampling from `context`.
|
|
"""
|
|
logging.info("Estimating loglikelihood for %d pairs." % len(requests))
|
|
|
|
group = mx.distributed.init()
|
|
|
|
# Group by common prefix
|
|
group_reqs = collections.defaultdict(list)
|
|
for idx, req in enumerate(requests):
|
|
group_reqs[req.args[0]].append((idx, req.args[1]))
|
|
questions = list(group_reqs.keys())
|
|
responses = []
|
|
indices = []
|
|
for v in group_reqs.values():
|
|
idx, resp = zip(*v)
|
|
indices.append(idx)
|
|
responses.append(resp)
|
|
|
|
# split data accross ranks
|
|
questions = questions[group.rank() :: group.size()]
|
|
responses = responses[group.rank() :: group.size()]
|
|
|
|
long_completions = 0
|
|
scores, is_greedy = [], []
|
|
for q, rs in tqdm(zip(questions, responses), total=len(questions)):
|
|
prefix = self._tokenize([q])[0]
|
|
full_sequences = self._tokenize([q + r for r in rs])
|
|
max_completed_l = max(len(s) for s in full_sequences)
|
|
|
|
# compute truncation length
|
|
max_tokens = self._max_tokens or DEFAULT_MAX_TOKENS
|
|
truncation = max(0, max_completed_l - max_tokens - 1)
|
|
orig_prefix_l = len(prefix)
|
|
prefix_l = max(len(prefix) - truncation, 0)
|
|
prefix = prefix[len(prefix) - prefix_l :]
|
|
|
|
# If the entire prompt got truncated ignore the question
|
|
if prefix_l == 0:
|
|
long_completions += 1
|
|
all_scores.extend([-float("inf")] * len(rs))
|
|
all_is_greedy.extend([False] * len(rs))
|
|
continue
|
|
|
|
# model scoring, returns num_requests x (logp, is_greedy, length).
|
|
logprobs, cache = self._process_prompt(prefix)
|
|
max_idx = mx.argmax(logprobs).item()
|
|
|
|
for s in full_sequences:
|
|
inputs = s[len(prefix) :]
|
|
# The logprobs from the last token of the prompt are
|
|
# for the first input token
|
|
scores.append(logprobs[0, inputs[0]].item())
|
|
is_greedy.append((inputs[0] == max_idx))
|
|
|
|
if len(inputs) == 1:
|
|
continue
|
|
score, _, ig = self._score_fn(
|
|
mx.array(inputs)[None, :], cache=copy.deepcopy(cache)
|
|
)
|
|
scores[-1] += mx.sum(score).item()
|
|
is_greedy[-1] &= mx.all(ig).item()
|
|
|
|
if long_completions > 0:
|
|
logging.info(
|
|
f"Prefix eliminated for {long_completions} requests with "
|
|
+ "completion longer than context."
|
|
)
|
|
|
|
# All gather the results across nodes
|
|
num_results = len(requests)
|
|
per_group = mx.distributed.all_max(len(scores), stream=mx.cpu).item()
|
|
scores = scores + [0] * (per_group - len(scores))
|
|
is_greedy = is_greedy + [False] * (per_group - len(is_greedy))
|
|
scores = mx.array(scores)
|
|
is_greedy = mx.array(is_greedy)
|
|
scores = mx.distributed.all_gather(scores, stream=mx.cpu)
|
|
is_greedy = mx.distributed.all_gather(is_greedy, stream=mx.cpu)
|
|
mx.eval(scores, is_greedy)
|
|
|
|
# Arrange the indices to match the scores from each node and then
|
|
# inverse sort the scores
|
|
all_indices = []
|
|
for rank in range(group.size()):
|
|
rank_indices = [
|
|
idx for question in indices[rank :: group.size()] for idx in question
|
|
]
|
|
rank_indices += [num_results] * (per_group - len(rank_indices))
|
|
all_indices.extend(rank_indices)
|
|
inv_sort = mx.argsort(mx.array(all_indices))
|
|
scores = scores[:num_results][inv_sort]
|
|
is_greedy = is_greedy[:num_results][inv_sort]
|
|
|
|
return list(zip(scores.tolist(), is_greedy.tolist()))
|
|
|
|
def loglikelihood_rolling(self, requests) -> list[float]:
|
|
"""Compute full log-likelihood of a string, with no truncation, for perplexity computation
|
|
- We will use the full max context length of the model.
|
|
- For inputs that exceed the max context length, we divide the tokenized string into chunks of up to
|
|
the max context length.
|
|
- IMPORTANT: Each document's loglikelihood/perplexity is computed *separately*, unlike other implementations
|
|
which may simply concatenate multiple documents together.
|
|
- IMPORTANT: We maximize the amount of context for each prediction. Specifically, for inputs that we break into
|
|
multiple chunks, the last input will still a full-sized context.
|
|
Example:
|
|
Input tokens: [ 0 1 2 3 4 5 6 7 8 9 ]
|
|
Prefix: EOT
|
|
Max context length: 4
|
|
Resulting input/prediction pairs:
|
|
INPUT: EOT 0 1 2
|
|
PRED: 0 1 2 3
|
|
INPUT: 3 4 5 6
|
|
PRED: 4 5 6 7
|
|
INPUT: 5 6 7 8
|
|
PRED: 8 9
|
|
Observe that:
|
|
1. Each token is predicted exactly once
|
|
2. For the last pair, we provide the full context, but only score the last two tokens
|
|
:param requests: list[Instance]
|
|
A list of Instance objects with property `args` which returns a tuple (context,).
|
|
string: str
|
|
String for which we are computing overall loglikelihood
|
|
:return: list[tuple[float]]
|
|
A list of tuples (logprob,)
|
|
logprob: float
|
|
The log probability of `context` conditioned on the EOT token.
|
|
"""
|
|
logging.info(
|
|
"Estimating loglikelihood rolling for %d sequences." % len(requests)
|
|
)
|
|
inputs = self._tokenize([req.args[0] for req in requests])
|
|
all_scores = []
|
|
for i in tqdm(range(0, len(inputs), self._batch_size)):
|
|
batch = inputs[i : i + self._batch_size]
|
|
scores, lengths, _ = self._score_fn(batch)
|
|
mask = mx.arange(scores.shape[-1]) < lengths[:, None]
|
|
all_scores.extend((mask * scores).sum(axis=-1).tolist())
|
|
|
|
return all_scores
|
|
|
|
def generate_until(self, requests) -> list[str]:
|
|
"""Generate greedily until a stopping sequence
|
|
:param requests: list[Instance]
|
|
A list of Instance objects with property `args` which returns a tuple (context, until).
|
|
context: str
|
|
Context string
|
|
until: [str]
|
|
The string sequences to generate until. These string sequences
|
|
may each span across multiple tokens, or may be part of one token.
|
|
:return: list[str]
|
|
A list of strings continuation
|
|
continuation: str
|
|
The generated continuation.
|
|
"""
|
|
group = mx.distributed.init()
|
|
|
|
# split data accross ranks
|
|
total_requests = len(requests)
|
|
requests = requests[group.rank() :: group.size()]
|
|
|
|
logging.info("Generating continuation for %d sequences." % len(requests))
|
|
contexts, options = zip(*[req.args for req in requests])
|
|
# The second element of the tuple contains:
|
|
# {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
|
|
|
|
# Tokenize all contexts
|
|
contexts = [
|
|
self.tokenizer.encode(
|
|
context, add_special_tokens=not self.use_chat_template
|
|
)
|
|
for context in contexts
|
|
]
|
|
|
|
# TODO consider multi-token, per-prompt stop conditions
|
|
max_tokens = [
|
|
self._max_tokens or opt.get("max_gen_tokens", DEFAULT_MAX_TOKENS)
|
|
for opt in options
|
|
]
|
|
|
|
completions = batch_generate(
|
|
model=self._model,
|
|
tokenizer=self.tokenizer,
|
|
prompts=contexts,
|
|
max_tokens=max_tokens,
|
|
verbose=True,
|
|
sampler=self._sampler,
|
|
).texts
|
|
|
|
for e, (text, opt) in enumerate(zip(completions, options)):
|
|
completions[e] = _rstrip_until(text, opt["until"])
|
|
if self.tokenizer.has_thinking:
|
|
completions[e] = _lstrip(text, self.tokenizer.think_end)
|
|
|
|
# Gather the completions
|
|
if group.size() > 1:
|
|
with mx.stream(mx.cpu):
|
|
pad_to = (total_requests + group.size() - 1) // group.size()
|
|
pad = pad_to - len(completions)
|
|
completions = [list(c.encode("utf-8")) for c in completions]
|
|
max_len = mx.array(max(len(c) for c in completions))
|
|
max_len = mx.distributed.all_max(max_len).item()
|
|
lengths = mx.array([len(c) for c in completions] + [0] * pad)
|
|
completions = mx.array(
|
|
[c + [0] * (max_len - len(c)) for c in completions]
|
|
+ [[0] * max_len] * pad,
|
|
mx.uint8,
|
|
)
|
|
completions = (
|
|
mx.distributed.all_gather(completions[None])
|
|
.swapaxes(0, 1)
|
|
.flatten(0, 1)
|
|
.tolist()
|
|
)
|
|
lengths = (
|
|
mx.distributed.all_gather(lengths[None])
|
|
.swapaxes(0, 1)
|
|
.flatten(0, 1)
|
|
.tolist()
|
|
)
|
|
completions = completions[:total_requests]
|
|
lengths = lengths[:total_requests]
|
|
completions = [
|
|
bytearray(c[:l]).decode() for c, l in zip(completions, lengths)
|
|
]
|
|
|
|
return completions
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
"Evaluate an MLX model using lm-evaluation-harness."
|
|
)
|
|
parser.add_argument("--model", help="Model to evaluate", required=True)
|
|
parser.add_argument("--tasks", nargs="+", required=True)
|
|
parser.add_argument(
|
|
"--output-dir", default=".", help="Output directory for result files."
|
|
)
|
|
parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
|
|
parser.add_argument("--num-shots", type=int, default=None, help="Number of shots")
|
|
parser.add_argument(
|
|
"--max-tokens",
|
|
type=int,
|
|
help="Maximum number of tokens to generate. When set, this value takes"
|
|
" precedence over task specific defaults.",
|
|
default=None,
|
|
)
|
|
parser.add_argument(
|
|
"--limit",
|
|
default=None,
|
|
help="Limit the number of examples per task.",
|
|
type=int,
|
|
)
|
|
parser.add_argument("--seed", type=int, default=123, help="Random seed.")
|
|
parser.add_argument(
|
|
"--fewshot-as-multiturn",
|
|
action="store_true",
|
|
help="Whether to provide the fewshot examples as a multiturn "
|
|
"conversation or a single user turn.",
|
|
default=False,
|
|
)
|
|
parser.add_argument(
|
|
"--apply-chat-template",
|
|
action=argparse.BooleanOptionalAction,
|
|
help="Specifies whether to apply a chat template to the prompt. If "
|
|
"the model has a chat template, this defaults to `True`, "
|
|
"otherwise `False`.",
|
|
default=None,
|
|
)
|
|
parser.add_argument(
|
|
"--chat-template-args",
|
|
type=json.loads,
|
|
help="""A JSON formatted string of arguments for the tokenizer's
|
|
apply_chat_template, e.g. '{"enable_thinking":false}'""",
|
|
default="{}",
|
|
)
|
|
parser.add_argument(
|
|
"--confirm-run-unsafe-code",
|
|
action="store_true",
|
|
help="Confirm that you want to run tasks that execute untrusted code.",
|
|
default=False,
|
|
)
|
|
parser.add_argument(
|
|
"--trust-remote-code",
|
|
action="store_true",
|
|
help="Enable trusting remote code for tokenizer",
|
|
)
|
|
parser.add_argument("--temp", type=float, default=0.0, help="Sampling temperature")
|
|
parser.add_argument("--top-p", type=float, default=1.0, help="Sampling top-p")
|
|
parser.add_argument("--top-k", type=int, default=0, help="Sampling top-k")
|
|
args = parser.parse_args()
|
|
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Silence tokenizer warnings
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
mx.random.seed(args.seed)
|
|
|
|
# Initialize the communication if in distributed mode
|
|
world = mx.distributed.init()
|
|
mx.eval(mx.distributed.all_sum(1, stream=mx.cpu))
|
|
if world.size() > 1 and world.rank() == 0:
|
|
print(f"Evaluating with {world.size()} nodes")
|
|
|
|
sampler = make_sampler(
|
|
temp=args.temp,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
)
|
|
lm = MLXLM(
|
|
args.model,
|
|
max_tokens=args.max_tokens,
|
|
batch_size=args.batch_size,
|
|
use_chat_template=args.apply_chat_template,
|
|
trust_remote_code=args.trust_remote_code,
|
|
sampler=sampler,
|
|
)
|
|
MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
|
|
|
|
results = lm_eval.simple_evaluate(
|
|
model=lm,
|
|
tasks=args.tasks,
|
|
fewshot_as_multiturn=args.fewshot_as_multiturn,
|
|
apply_chat_template=lm.use_chat_template,
|
|
num_fewshot=args.num_shots,
|
|
limit=args.limit,
|
|
random_seed=args.seed,
|
|
numpy_random_seed=args.seed,
|
|
torch_random_seed=args.seed,
|
|
fewshot_random_seed=args.seed,
|
|
confirm_run_unsafe_code=args.confirm_run_unsafe_code,
|
|
)
|
|
|
|
file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
|
|
if args.num_shots is not None:
|
|
file_keys += [f"{args.num_shots:02d}"]
|
|
file_keys += args.tasks
|
|
filename = "_".join(file_keys)
|
|
if world.rank() == 0:
|
|
output_path = output_dir / filename
|
|
output_path.write_text(json.dumps(results["results"], indent=4))
|
|
print("Results:")
|
|
for result in results["results"].values():
|
|
print(json.dumps(result, indent=4))
|