75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Run one MiniMind-3 inference.")
|
|
parser.add_argument("--model_path", default=str(ROOT / "minimind-3"))
|
|
parser.add_argument("--prompt", default="MiniMind 是什么？")
|
|
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
parser.add_argument("--max_new_tokens", type=int, default=64)
|
|
parser.add_argument("--open_thinking", action="store_true")
|
|
parser.add_argument("--do_sample", action="store_true")
|
|
parser.add_argument("--temperature", type=float, default=0.85)
|
|
parser.add_argument("--top_p", type=float, default=0.95)
|
|
args = parser.parse_args()
|
|
|
|
model_path = Path(args.model_path)
|
|
if not model_path.exists():
|
|
raise FileNotFoundError(f"model path not found: {model_path}")
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(str(model_path))
|
|
model = AutoModelForCausalLM.from_pretrained(str(model_path), trust_remote_code=True)
|
|
model = model.eval().to(args.device)
|
|
|
|
messages = [{"role": "user", "content": args.prompt}]
|
|
prompt_text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True,
|
|
open_thinking=args.open_thinking,
|
|
)
|
|
inputs = tokenizer(prompt_text, return_tensors="pt").to(args.device)
|
|
|
|
generate_kwargs = {
|
|
"inputs": inputs["input_ids"],
|
|
"attention_mask": inputs["attention_mask"],
|
|
"max_new_tokens": args.max_new_tokens,
|
|
"do_sample": args.do_sample,
|
|
"pad_token_id": tokenizer.pad_token_id,
|
|
"eos_token_id": tokenizer.eos_token_id,
|
|
}
|
|
if args.do_sample:
|
|
generate_kwargs["temperature"] = args.temperature
|
|
generate_kwargs["top_p"] = args.top_p
|
|
|
|
with torch.no_grad():
|
|
generated_ids = model.generate(**generate_kwargs)
|
|
|
|
prompt_len = inputs["input_ids"].shape[1]
|
|
response_ids = generated_ids[0, prompt_len:]
|
|
response = tokenizer.decode(response_ids, skip_special_tokens=True)
|
|
|
|
print("[Prompt]")
|
|
print(args.prompt)
|
|
print()
|
|
print("[Shapes]")
|
|
print(f"input_ids.shape={tuple(inputs['input_ids'].shape)}")
|
|
print(f"generated_ids.shape={tuple(generated_ids.shape)}")
|
|
print(f"new_tokens={response_ids.numel()}")
|
|
print()
|
|
print("[Response]")
|
|
print(response.strip())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
main()
|