Spaces:
Running
Running
File size: 4,335 Bytes
a61f6ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any
from model_release.blum_finance.inference import _extract_json_object
from model_release.evaluation.evaluate_candidate import load_examples
from model_release.evaluation.evaluate_mlx_candidate import write_evaluation
def generate_predictions(
*,
model_path: Path,
examples: list[dict[str, Any]],
max_tokens: int,
device: str,
batch_size: int,
) -> tuple[list[dict[str, Any] | None], list[dict[str, Any]]]:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
model_path,
dtype=torch.bfloat16,
low_cpu_mem_usage=True,
).to(device)
tokenizer.padding_side = "left"
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
predictions: list[dict[str, Any] | None] = []
generations: list[dict[str, Any]] = []
for batch_start in range(0, len(examples), batch_size):
batch = examples[batch_start : batch_start + batch_size]
prompts = [
tokenizer.apply_chat_template(
example["messages"][:-1],
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
)
for example in batch
]
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(device)
with torch.inference_mode():
output = model.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=False,
pad_token_id=tokenizer.eos_token_id,
)
generated_rows = tokenizer.batch_decode(
output[:, inputs.input_ids.shape[1] :],
skip_special_tokens=True,
)
errors: list[str | None] = []
for batch_index, (example, generated) in enumerate(
zip(batch, generated_rows, strict=True)
):
index = batch_start + batch_index
try:
parsed = _extract_json_object(generated)
error = None
except (ValueError, json.JSONDecodeError) as exc:
parsed = None
error = str(exc)
predictions.append(parsed)
generations.append(
{
"example_id": example.get("example_id"),
"index": index,
"generated_text": generated,
"parse_error": error,
}
)
errors.append(error)
print(
json.dumps(
{
"completed": batch_start + len(batch),
"total": len(examples),
"parse_errors": sum(error is not None for error in errors),
}
),
flush=True,
)
return predictions, generations
def main() -> None:
parser = argparse.ArgumentParser(
description="Evaluate a local merged Transformers BLUM Finance candidate."
)
parser.add_argument("--model-path", type=Path, required=True)
parser.add_argument("--revision", required=True)
parser.add_argument("--test-file", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--max-tokens", type=int, default=512)
parser.add_argument("--device", default="mps")
parser.add_argument("--batch-size", type=int, default=1)
parser.add_argument("--limit", type=int)
args = parser.parse_args()
examples = load_examples(args.test_file)
if args.limit:
examples = examples[: args.limit]
predictions, generations = generate_predictions(
model_path=args.model_path,
examples=examples,
max_tokens=args.max_tokens,
device=args.device,
batch_size=args.batch_size,
)
payload = write_evaluation(
output_dir=args.output_dir,
model_revision=args.revision,
examples=examples,
predictions=predictions,
generations=generations,
)
print(json.dumps(payload, indent=2, sort_keys=True))
if __name__ == "__main__":
main()
|