| |
| """Run PertMind with Hugging Face Transformers.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
| DEFAULT_SYSTEM_PROMPT = ( |
| "You are PertMind, a biomedical assistant. For biomedical prediction, " |
| "screen-ranking, or gene-set interpretation tasks, answer first and then " |
| "provide a concise explanation. Use this style when applicable:\n" |
| "Final Answer: <answer>\nExplanation: <brief explanation>" |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model", default=".", help="Path or Hugging Face model id.") |
| parser.add_argument("--prompt", required=True, help="User prompt.") |
| parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT) |
| parser.add_argument("--max-new-tokens", type=int, default=768) |
| parser.add_argument("--temperature", type=float, default=0.0) |
| parser.add_argument("--top-p", type=float, default=0.95) |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| args.model, |
| torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, |
| device_map="auto", |
| trust_remote_code=True, |
| ) |
| messages = [ |
| {"role": "system", "content": args.system_prompt}, |
| {"role": "user", "content": args.prompt}, |
| ] |
| try: |
| text = tokenizer.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| except TypeError: |
| text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = tokenizer(text, return_tensors="pt").to(model.device) |
| do_sample = args.temperature > 0 |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=args.max_new_tokens, |
| do_sample=do_sample, |
| temperature=args.temperature if do_sample else None, |
| top_p=args.top_p if do_sample else None, |
| pad_token_id=tokenizer.eos_token_id, |
| ) |
| generated = outputs[0, inputs["input_ids"].shape[-1] :] |
| print(tokenizer.decode(generated, skip_special_tokens=True).strip()) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
