fix: eval script bugs and add missing features
- evaluate_mmlu: fix double few-shot injection (build_prompt no longer
adds few-shot, apply_chat handles it once)
- evaluate_humaneval: fix pass@k k-filtering to be per-problem instead
of using first problem's n globally; reuse ProcessPoolExecutor across
problems; fix closure UnboundLocalError in test_one; handle None in
report when k > n
- evaluate_ifd: remove dead code (score_plain/score_messages); add
multi-file/directory input support with --input_path/--output_dir;
add summary.json aggregation and --max_samples; add --dtype flag
- evaluate_ppl: add --device and --dtype flags (was hardcoded to cuda)
- evaluate_ifeval: fix docstring path (scripts/tools -> scripts/eval)
- analyze_weights: add --output JSON export; fix dead code filter
("_norm" not in r was always True)
This commit is contained in:
parent
b12b24eadc
commit
c17aa0dc54
|
|
@ -117,7 +117,7 @@ def print_component_summary(results: dict[str, dict], title: str):
|
|||
r["er_99_norm"]
|
||||
for vs in matrix_groups.values()
|
||||
for r in vs
|
||||
if "_norm" not in r or not r.get("is_1d")
|
||||
if not r.get("is_1d")
|
||||
]
|
||||
if all_er:
|
||||
m = sum(all_er) / len(all_er)
|
||||
|
|
@ -232,8 +232,16 @@ def main():
|
|||
action="store_true",
|
||||
help="Skip SVD analysis, only show weight statistics (mean/std/min/max).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Save results as JSON to this path.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
all_results = {}
|
||||
|
||||
def analyze_one(ckpt_dir: str, label: str):
|
||||
ckpt_dir = Path(ckpt_dir)
|
||||
weights_path = ckpt_dir / "model.safetensors"
|
||||
|
|
@ -294,13 +302,19 @@ def main():
|
|||
)
|
||||
print_layer_grid(results)
|
||||
print_weight_stats(results)
|
||||
all_results[label] = results
|
||||
return results
|
||||
|
||||
analyze_one(args.ckpt_dir, "Primary")
|
||||
|
||||
if args.compare:
|
||||
for cdir in args.compare:
|
||||
analyze_one(cdir, "Compare")
|
||||
analyze_one(cdir, f"Compare_{cdir}")
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"\nResults saved to {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ Config is a single dataclass; side effects are isolated at pipeline boundaries.
|
|||
"""
|
||||
|
||||
import argparse
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
|
|
@ -233,7 +232,7 @@ def execute_one(args: tuple) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]:
|
||||
def test_one(item: dict, cfg: EvalConfig, pool=None) -> Tuple[str, int, int]:
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
|
||||
task_id = item["task_id"]
|
||||
|
|
@ -247,11 +246,16 @@ def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]:
|
|||
for c in completions
|
||||
]
|
||||
n = len(codes)
|
||||
passed = 0
|
||||
with ProcessPoolExecutor(max_workers=cfg.test_workers) as pool:
|
||||
for ok in pool.map(execute_one, codes):
|
||||
if ok:
|
||||
passed += 1
|
||||
|
||||
def _run(p):
|
||||
return sum(1 for ok in p.map(execute_one, codes) if ok)
|
||||
|
||||
if pool is not None:
|
||||
passed = _run(pool)
|
||||
else:
|
||||
with ProcessPoolExecutor(max_workers=cfg.test_workers) as p:
|
||||
passed = _run(p)
|
||||
|
||||
return task_id, n, passed
|
||||
|
||||
|
||||
|
|
@ -259,8 +263,14 @@ def test_all(
|
|||
items: Sequence[dict],
|
||||
cfg: EvalConfig,
|
||||
) -> Iterator[Tuple[str, int, int]]:
|
||||
for item in tqdm.tqdm(items, desc="Testing", unit="problem"):
|
||||
yield test_one(item, cfg)
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
|
||||
pool = ProcessPoolExecutor(max_workers=cfg.test_workers)
|
||||
try:
|
||||
for item in tqdm.tqdm(items, desc="Testing", unit="problem"):
|
||||
yield test_one(item, cfg, pool)
|
||||
finally:
|
||||
pool.shutdown(wait=True)
|
||||
|
||||
|
||||
def pass_at_k(n: int, c: int, k: int) -> float:
|
||||
|
|
@ -273,26 +283,32 @@ def score_results(
|
|||
results: Iterator[Tuple[str, int, int]],
|
||||
k_values: Tuple[int, ...],
|
||||
) -> Dict:
|
||||
# filter to k <= n (peek first result to get n)
|
||||
first = next(results)
|
||||
results = itertools.chain([first], results)
|
||||
n = first[1]
|
||||
k_values = tuple(k for k in k_values if k <= n)
|
||||
"""Score pass@k for each problem.
|
||||
|
||||
k values are filtered per-problem: if a problem has n < k samples
|
||||
(e.g. after deduplication), pass@k is not computed for that problem.
|
||||
The summary averages only over problems where the k was computed.
|
||||
"""
|
||||
scores = {k: [] for k in k_values}
|
||||
output = {}
|
||||
for task_id, n, passed in results:
|
||||
entry = {"task_id": task_id, "n": n, "passed": passed}
|
||||
for k in k_values:
|
||||
pk = round(pass_at_k(n, passed, k), 4)
|
||||
entry[f"pass@{k}"] = pk
|
||||
scores[k].append(pk)
|
||||
if k <= n:
|
||||
pk = round(pass_at_k(n, passed, k), 4)
|
||||
entry[f"pass@{k}"] = pk
|
||||
scores[k].append(pk)
|
||||
else:
|
||||
entry[f"pass@{k}"] = None
|
||||
output[task_id] = entry
|
||||
|
||||
summary = {}
|
||||
for k in k_values:
|
||||
vals = scores[k]
|
||||
summary[f"pass@{k}"] = round(float(np.mean(vals)), 4)
|
||||
if vals:
|
||||
summary[f"pass@{k}"] = round(float(np.mean(vals)), 4)
|
||||
else:
|
||||
summary[f"pass@{k}"] = None
|
||||
output["_summary"] = summary
|
||||
return output
|
||||
|
||||
|
|
@ -375,7 +391,10 @@ def report(scored: Dict):
|
|||
summary = scored.pop("_summary", {})
|
||||
print(f"\n{'=' * 60}")
|
||||
for k, v in summary.items():
|
||||
print(f" {k}: {v:.2%}")
|
||||
if v is not None:
|
||||
print(f" {k}: {v:.2%}")
|
||||
else:
|
||||
print(f" {k}: N/A")
|
||||
print(f"{'=' * 60}")
|
||||
scored["_summary"] = summary
|
||||
|
||||
|
|
|
|||
|
|
@ -16,7 +16,9 @@ v2 changelog:
|
|||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
|
||||
import torch
|
||||
|
|
@ -65,6 +67,29 @@ def _resolve_sentinel_ids(tokenizer, sentinel_text):
|
|||
return [0]
|
||||
|
||||
|
||||
def _collect_input_files(input_path: str) -> list:
|
||||
"""Resolve *input_path* to a list of JSONL/JSON files."""
|
||||
if os.path.isdir(input_path):
|
||||
files = []
|
||||
for ext in ("*.jsonl", "*.json"):
|
||||
files.extend(
|
||||
sorted(glob.glob(os.path.join(input_path, "**", ext), recursive=True))
|
||||
)
|
||||
return files
|
||||
return sorted(glob.glob(input_path))
|
||||
|
||||
|
||||
def _load_items(filepath: str) -> list:
|
||||
"""Load JSONL or JSON (array / single dict) into a list of dicts."""
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
if filepath.lower().endswith(".json"):
|
||||
data = json.load(f)
|
||||
if isinstance(data, dict):
|
||||
return [data]
|
||||
return data
|
||||
return [json.loads(line) for line in f if line.strip()]
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _score_batch(
|
||||
pairs, model, device, max_len=2048, sentinel_ids=None, per_token=False
|
||||
|
|
@ -204,69 +229,9 @@ def _trim(context_ids, resp_ids, max_len):
|
|||
return context_ids[overflow:], resp_ids
|
||||
|
||||
|
||||
def score_plain(
|
||||
def process_file(
|
||||
model,
|
||||
tokenizer,
|
||||
instruction,
|
||||
response,
|
||||
device,
|
||||
max_len=2048,
|
||||
sentinel_ids=None,
|
||||
per_token=False,
|
||||
):
|
||||
"""Compute IFD for a single instruction-response pair (plain format)."""
|
||||
ctx_ids = tokenizer.encode(instruction, add_special_tokens=False)
|
||||
resp_ids = tokenizer.encode(response, add_special_tokens=False)
|
||||
ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
|
||||
if not ctx_ids or not resp_ids:
|
||||
return {
|
||||
"L_cond": None,
|
||||
"L_uncond": None,
|
||||
"ifd": None,
|
||||
"skip_reason": "empty ctx or resp",
|
||||
}
|
||||
return _score_batch(
|
||||
[(ctx_ids, resp_ids)],
|
||||
model,
|
||||
device,
|
||||
max_len,
|
||||
sentinel_ids=sentinel_ids,
|
||||
per_token=per_token,
|
||||
)[0]
|
||||
|
||||
|
||||
def score_messages(
|
||||
model, tokenizer, messages, device, max_len=2048, sentinel_ids=None, per_token=False
|
||||
):
|
||||
"""Compute IFD for each assistant turn in a messages array."""
|
||||
turns = []
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.get("role") != "assistant":
|
||||
continue
|
||||
ctx_text = "\n\n".join(m["content"] for m in messages[:i])
|
||||
ctx_ids = tokenizer.encode(ctx_text)
|
||||
resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
|
||||
ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
|
||||
if ctx_ids and resp_ids:
|
||||
turns.append((ctx_ids, resp_ids))
|
||||
if not turns:
|
||||
return None
|
||||
raw_scores = _score_batch(
|
||||
turns, model, device, max_len, sentinel_ids=sentinel_ids, per_token=per_token
|
||||
)
|
||||
valid = [s for s in raw_scores if s is not None and s.get("ifd") is not None]
|
||||
if not valid:
|
||||
return {"ifd": None, "ifd_turns": raw_scores}
|
||||
avg = sum(s["ifd"] for s in valid) / len(valid)
|
||||
return {
|
||||
"ifd": avg,
|
||||
"ifd_detail": valid[0] if len(valid) == 1 else None,
|
||||
"ifd_turns": raw_scores,
|
||||
}
|
||||
|
||||
|
||||
def process_file(
|
||||
param_path,
|
||||
input_file,
|
||||
output_file,
|
||||
instr_key,
|
||||
|
|
@ -275,28 +240,31 @@ def process_file(
|
|||
data_format="plain",
|
||||
batch_size=1,
|
||||
device=None,
|
||||
sentinel_text="\n",
|
||||
sentinel_ids=None,
|
||||
per_token=False,
|
||||
max_samples=None,
|
||||
):
|
||||
"""Score a single file, write per-sample JSONL, return summary stats."""
|
||||
if device is None:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16 if "cuda" in device else torch.float32
|
||||
|
||||
model = AutoModel.from_pretrained(param_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(param_path)
|
||||
model.to(device=device, dtype=dtype)
|
||||
model.eval()
|
||||
if sentinel_ids is None:
|
||||
sentinel_ids = _resolve_sentinel_ids(tokenizer, "\n")
|
||||
|
||||
sentinel_ids = _resolve_sentinel_ids(tokenizer, sentinel_text)
|
||||
data = _load_items(input_file)
|
||||
|
||||
with open(input_file, encoding="utf-8") as f:
|
||||
data = [json.loads(line) for line in f if line.strip()]
|
||||
if max_samples and len(data) > max_samples:
|
||||
import random
|
||||
|
||||
data = random.sample(data, max_samples)
|
||||
|
||||
results = []
|
||||
all_ifds = []
|
||||
buffer = []
|
||||
|
||||
for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
|
||||
label = os.path.splitext(os.path.basename(input_file))[0]
|
||||
|
||||
for item in tqdm.tqdm(data, desc=f" {label}", unit="sample", leave=False):
|
||||
if data_format == "messages":
|
||||
turns = []
|
||||
for i, msg in enumerate(item.get("messages", [])):
|
||||
|
|
@ -356,8 +324,22 @@ def process_file(
|
|||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
|
||||
valid_ifd = [v for v in all_ifds if v is not None]
|
||||
stats = {
|
||||
"samples": len(data),
|
||||
"valid_ifd": len(valid_ifd),
|
||||
"skipped": len(data) - len(valid_ifd),
|
||||
}
|
||||
if valid_ifd:
|
||||
stats["mean_ifd"] = statistics.mean(valid_ifd)
|
||||
stats["median_ifd"] = statistics.median(valid_ifd)
|
||||
if len(valid_ifd) > 1:
|
||||
stats["stdev_ifd"] = statistics.stdev(valid_ifd)
|
||||
stats["min_ifd"] = min(valid_ifd)
|
||||
stats["max_ifd"] = max(valid_ifd)
|
||||
|
||||
print(f"\n{'=' * 50}")
|
||||
print(f" [{label}]")
|
||||
print(f"{'=' * 50}")
|
||||
print(f" Samples: {len(data)}")
|
||||
print(f" Valid IFD: {len(valid_ifd)}")
|
||||
print(f" Skipped: {len(data) - len(valid_ifd)}")
|
||||
|
|
@ -368,7 +350,8 @@ def process_file(
|
|||
print(f" Min IFD: {min(valid_ifd):.4f}")
|
||||
print(f" Max IFD: {max(valid_ifd):.4f}")
|
||||
print(f"{'=' * 50}")
|
||||
print(f"Results saved to {output_file}")
|
||||
print(f" Results saved to {output_file}")
|
||||
return stats
|
||||
|
||||
|
||||
def _flush_buffer(
|
||||
|
|
@ -422,8 +405,18 @@ def main():
|
|||
description="Compute IFD scores for instruction-response data"
|
||||
)
|
||||
parser.add_argument("--param_path", type=str, required=True, help="Model directory")
|
||||
parser.add_argument("--input", type=str, required=True, help="Input JSONL file")
|
||||
parser.add_argument("--output", type=str, required=True, help="Output JSONL file")
|
||||
parser.add_argument(
|
||||
"--input_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Input file, glob pattern, or directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory for output files (summary.json + per-file JSONL).",
|
||||
)
|
||||
parser.add_argument("--max_len", type=int, default=2048, help="Max token length")
|
||||
parser.add_argument(
|
||||
"--format",
|
||||
|
|
@ -442,6 +435,12 @@ def main():
|
|||
"--batch_size", type=int, default=8, help="Batch size for model forward passes"
|
||||
)
|
||||
parser.add_argument("--device", type=str, default=None, help="Device (e.g. cuda:0)")
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="bfloat16" if torch.cuda.is_available() else "float32",
|
||||
help="Torch dtype",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sentinel_text",
|
||||
type=str,
|
||||
|
|
@ -453,21 +452,60 @@ def main():
|
|||
action="store_true",
|
||||
help="Include per-token IFD breakdown in output",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_samples",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum number of samples per file (random subsample). Default: all.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
process_file(
|
||||
args.param_path,
|
||||
args.input,
|
||||
args.output,
|
||||
args.instr_key,
|
||||
args.resp_key,
|
||||
args.max_len,
|
||||
data_format=args.format,
|
||||
batch_size=args.batch_size,
|
||||
device=args.device,
|
||||
sentinel_text=args.sentinel_text,
|
||||
per_token=args.per_token,
|
||||
)
|
||||
if args.device is None:
|
||||
args.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = getattr(torch, args.dtype)
|
||||
|
||||
print(f"Loading model from {args.param_path} ...")
|
||||
model = AutoModel.from_pretrained(args.param_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.param_path)
|
||||
model.to(device=args.device, dtype=dtype)
|
||||
model.eval()
|
||||
|
||||
sentinel_ids = _resolve_sentinel_ids(tokenizer, args.sentinel_text)
|
||||
|
||||
input_files = _collect_input_files(args.input_path)
|
||||
if not input_files:
|
||||
print(f"No input files found at {args.input_path}")
|
||||
return
|
||||
|
||||
print(f"Found {len(input_files)} file(s) to evaluate")
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
all_stats = {}
|
||||
for filepath in input_files:
|
||||
label = os.path.splitext(os.path.basename(filepath))[0]
|
||||
output_file = os.path.join(args.output_dir, f"{label}_ifd.jsonl")
|
||||
|
||||
stats = process_file(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
input_file=filepath,
|
||||
output_file=output_file,
|
||||
instr_key=args.instr_key,
|
||||
resp_key=args.resp_key,
|
||||
max_len=args.max_len,
|
||||
data_format=args.format,
|
||||
batch_size=args.batch_size,
|
||||
device=args.device,
|
||||
sentinel_ids=sentinel_ids,
|
||||
per_token=args.per_token,
|
||||
max_samples=args.max_samples,
|
||||
)
|
||||
all_stats[label] = stats
|
||||
|
||||
summary_path = os.path.join(args.output_dir, "summary.json")
|
||||
with open(summary_path, "w", encoding="utf-8") as f:
|
||||
json.dump(all_stats, f, ensure_ascii=False, indent=2)
|
||||
print(f"\nSummary saved to {summary_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Supports all IFEval constraint types except language detection.
|
|||
|
||||
Usage::
|
||||
|
||||
python scripts/tools/evaluate_ifeval.py --param_path ./params \
|
||||
python scripts/eval/evaluate_ifeval.py --param_path ./params \
|
||||
--data_path ifeval.jsonl --output results.json \
|
||||
--temperature 0.1 --max_tokens 512
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -139,17 +139,12 @@ def load_csv(path: str) -> list[dict]:
|
|||
return data
|
||||
|
||||
|
||||
def build_prompt(
|
||||
question: str, choices: dict, subject: str, n_shot: int, dev_data: list[dict]
|
||||
) -> str:
|
||||
prompt = ""
|
||||
if n_shot > 0 and dev_data:
|
||||
prompt = f"The following are multiple choice questions (with answers) about {subject}.\n\n"
|
||||
for item in dev_data[:n_shot]:
|
||||
prompt += f"Question: {item['question']}\n"
|
||||
for k in ("A", "B", "C", "D"):
|
||||
prompt += f"{k}. {item[k]}\n"
|
||||
prompt += f"Answer: {item['answer']}\n\n"
|
||||
def build_prompt(question: str, choices: dict, subject: str) -> str:
|
||||
"""Build the raw question prompt (without few-shot examples).
|
||||
|
||||
Few-shot examples are handled by ``apply_chat`` to avoid duplication.
|
||||
"""
|
||||
prompt = f"The following are multiple choice questions (with answers) about {subject}.\n\n"
|
||||
prompt += f"Question: {question}\n"
|
||||
for k in ("A", "B", "C", "D"):
|
||||
prompt += f"{k}. {choices[k]}\n"
|
||||
|
|
@ -218,9 +213,7 @@ def evaluate_subject(
|
|||
correct = 0
|
||||
total = 0
|
||||
for item in tqdm.tqdm(test_data, desc=f"{subject:40s}", leave=False):
|
||||
raw_prompt = build_prompt(
|
||||
item["question"], item, subject, n_shot, dev_data or []
|
||||
)
|
||||
raw_prompt = build_prompt(item["question"], item, subject)
|
||||
context = apply_chat(tokenizer, raw_prompt, n_shot, dev_data or [])
|
||||
context_ids = tokenizer.encode(context)
|
||||
scores = {
|
||||
|
|
|
|||
|
|
@ -221,6 +221,7 @@ def process_file(
|
|||
max_samples: Optional[int],
|
||||
output_file: Optional[str],
|
||||
label: str,
|
||||
device: str = "cuda",
|
||||
) -> Dict:
|
||||
"""Evaluate a single dataset (list of items), return summary stats.
|
||||
|
||||
|
|
@ -252,8 +253,8 @@ def process_file(
|
|||
batch_texts = texts[i : i + batch_size]
|
||||
padded_ids, masks = _encode_batch(tokenizer, batch_texts, max_length)
|
||||
|
||||
input_ids = torch.tensor(padded_ids, device="cuda", dtype=torch.long)
|
||||
attention_mask = torch.tensor(masks, device="cuda", dtype=torch.bool)
|
||||
input_ids = torch.tensor(padded_ids, device=device, dtype=torch.long)
|
||||
attention_mask = torch.tensor(masks, device=device, dtype=torch.bool)
|
||||
|
||||
token_log_probs, valid = _compute_batch(model, input_ids, attention_mask)
|
||||
|
||||
|
|
@ -333,11 +334,14 @@ def main(
|
|||
max_length: int,
|
||||
token_level: bool,
|
||||
max_samples: Optional[int],
|
||||
device: str = "cuda",
|
||||
dtype: str = "bfloat16",
|
||||
):
|
||||
print(f"Loading model from {param_path} ...")
|
||||
model = AutoModel.from_pretrained(param_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(param_path)
|
||||
model.to(device="cuda", dtype=torch.bfloat16)
|
||||
torch_dtype = getattr(torch, dtype)
|
||||
model.to(device=device, dtype=torch_dtype)
|
||||
model.eval()
|
||||
|
||||
input_files = _collect_input_files(input_path)
|
||||
|
|
@ -371,6 +375,7 @@ def main(
|
|||
max_samples=max_samples,
|
||||
output_file=token_output,
|
||||
label=label,
|
||||
device=device,
|
||||
)
|
||||
all_stats[label] = stats
|
||||
print_stats(label, stats)
|
||||
|
|
@ -430,6 +435,18 @@ if __name__ == "__main__":
|
|||
default=None,
|
||||
help="Maximum number of samples per file (random subsample). Default: all.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device",
|
||||
type=str,
|
||||
default="cuda" if torch.cuda.is_available() else "cpu",
|
||||
help="Device for model inference.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="bfloat16" if torch.cuda.is_available() else "float32",
|
||||
help="Torch dtype for model weights.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
with torch.inference_mode():
|
||||
|
|
@ -442,4 +459,6 @@ if __name__ == "__main__":
|
|||
max_length=args.max_length,
|
||||
token_level=args.token_level,
|
||||
max_samples=args.max_samples,
|
||||
device=args.device,
|
||||
dtype=args.dtype,
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue