160 lines
6.7 KiB
Python
160 lines
6.7 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import torch
|
|
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from dataset.search_shortqa_dataset import build_search_shortqa_messages
|
|
from model.model_minimind import MiniMindForCausalLM
|
|
from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
|
|
from trainer.trainer_utils import (
|
|
add_model_profile_args,
|
|
apply_model_profile,
|
|
build_lm_config,
|
|
get_model_params,
|
|
get_model_suffix,
|
|
resolve_attention_type,
|
|
)
|
|
|
|
|
|
def iter_jsonl(path):
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
yield json.loads(line)
|
|
|
|
|
|
def load_model(args):
|
|
try:
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
except ImportError as exc:
|
|
raise ImportError("predict_search_shortqa.py requires transformers. Install project requirements before inference.") from exc
|
|
tokenizer = AutoTokenizer.from_pretrained(args.load_from)
|
|
if args.load_from == "model" or os.path.basename(args.load_from.rstrip("/")) == "model":
|
|
config = build_lm_config(
|
|
hidden_size=args.hidden_size,
|
|
num_hidden_layers=args.num_hidden_layers,
|
|
use_moe=bool(args.use_moe),
|
|
attention_type=resolve_attention_type(args),
|
|
kv_lora_rank=args.kv_lora_rank,
|
|
q_lora_rank=args.q_lora_rank,
|
|
rope_dim=args.rope_dim,
|
|
num_attention_heads=args.num_attention_heads,
|
|
num_key_value_heads=args.num_key_value_heads,
|
|
intermediate_size=args.intermediate_size,
|
|
)
|
|
model = MiniMindMLAForCausalLM(config) if isinstance(config, MiniMindMLAConfig) else MiniMindForCausalLM(config)
|
|
model_suffix = get_model_suffix(config)
|
|
ckp = os.path.join(args.save_dir, f"{args.weight}_{args.hidden_size}{model_suffix}.pth")
|
|
model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
|
|
else:
|
|
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
|
|
get_model_params(model, model.config)
|
|
model = model.eval().to(args.device)
|
|
if args.dtype == "float16" and "cuda" in args.device:
|
|
model = model.half()
|
|
elif args.dtype == "bfloat16" and "cuda" in args.device:
|
|
model = model.to(torch.bfloat16)
|
|
return model, tokenizer
|
|
|
|
|
|
@torch.inference_mode()
|
|
def generate_one(model, tokenizer, sample, args):
|
|
messages = build_search_shortqa_messages(
|
|
sample,
|
|
max_contexts=args.max_contexts,
|
|
max_chars_per_context=args.max_chars_per_context,
|
|
)
|
|
prompt = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True,
|
|
open_thinking=False,
|
|
)
|
|
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=args.max_seq_len).to(args.device)
|
|
start = time.time()
|
|
output_ids = model.generate(
|
|
inputs=inputs["input_ids"],
|
|
attention_mask=inputs.get("attention_mask"),
|
|
max_new_tokens=args.max_new_tokens,
|
|
do_sample=bool(args.do_sample),
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
repetition_penalty=args.repetition_penalty,
|
|
eos_token_id=tokenizer.eos_token_id,
|
|
)
|
|
elapsed = max(time.time() - start, 1e-6)
|
|
new_tokens = output_ids.size(1) - inputs["input_ids"].size(1)
|
|
response = tokenizer.decode(output_ids[0][inputs["input_ids"].size(1):], skip_special_tokens=True).strip()
|
|
return response, new_tokens / elapsed, new_tokens
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generate SearchShortQA prediction JSONL")
|
|
parser.add_argument("--data", required=True)
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--load_from", default="model")
|
|
parser.add_argument("--save_dir", default="out")
|
|
parser.add_argument("--weight", default="search_grpo")
|
|
parser.add_argument("--hidden_size", default=768, type=int)
|
|
parser.add_argument("--num_hidden_layers", default=8, type=int)
|
|
parser.add_argument("--num_attention_heads", default=8, type=int)
|
|
parser.add_argument("--num_key_value_heads", default=4, type=int)
|
|
parser.add_argument("--intermediate_size", default=None, type=int)
|
|
parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1])
|
|
parser.add_argument("--attention_type", default="gqa", choices=["gqa", "mha", "mqa", "mla"])
|
|
parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1])
|
|
parser.add_argument("--kv_lora_rank", default=128, type=int)
|
|
parser.add_argument("--q_lora_rank", default=256, type=int)
|
|
parser.add_argument("--rope_dim", default=None, type=int)
|
|
parser.add_argument("--max_seq_len", default=1536, type=int)
|
|
parser.add_argument("--max_new_tokens", default=128, type=int)
|
|
parser.add_argument("--max_contexts", default=6, type=int)
|
|
parser.add_argument("--max_chars_per_context", default=700, type=int)
|
|
parser.add_argument("--temperature", default=0.2, type=float)
|
|
parser.add_argument("--top_p", default=0.9, type=float)
|
|
parser.add_argument("--top_k", default=50, type=int)
|
|
parser.add_argument("--do_sample", default=0, type=int)
|
|
parser.add_argument("--repetition_penalty", default=1.05, type=float)
|
|
parser.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
|
|
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
parser.add_argument("--limit", default=0, type=int)
|
|
add_model_profile_args(parser)
|
|
args = parser.parse_args()
|
|
apply_model_profile(args)
|
|
|
|
model, tokenizer = load_model(args)
|
|
os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
|
|
total_tokens, total_tps, count = 0, 0.0, 0
|
|
with open(args.output, "w", encoding="utf-8") as out:
|
|
for sample in iter_jsonl(args.data):
|
|
if args.limit and count >= args.limit:
|
|
break
|
|
response, tps, new_tokens = generate_one(model, tokenizer, sample, args)
|
|
row = {
|
|
"id": sample.get("id"),
|
|
"question": sample.get("question"),
|
|
"prediction": response,
|
|
"tokens": new_tokens,
|
|
"tokens_per_second": tps,
|
|
}
|
|
out.write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
total_tokens += new_tokens
|
|
total_tps += tps
|
|
count += 1
|
|
if count % 20 == 0:
|
|
print(f"generated={count}, avg_tokens/s={total_tps / count:.2f}")
|
|
if count:
|
|
print(f"Wrote {count} predictions to {args.output}")
|
|
print(f"avg_tokens/s={total_tps / count:.2f}, total_new_tokens={total_tokens}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|