| import argparse |
| import torch |
| from pathlib import Path |
| from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM |
| from qwen4_exp_dense_compat import enable_dense_qwen4_exp_attention |
|
|
| p=argparse.ArgumentParser();p.add_argument('--raw',action='store_true');p.add_argument('--prompt',default='What is the capital of France? Answer in one sentence.');p.add_argument('--device',default='cuda:0');p.add_argument('--max-new-tokens',type=int,default=192);a=p.parse_args() |
| enable_dense_qwen4_exp_attention() |
| if a.device=='cpu': |
| from transformers.models.qwen4_exp import modeling_qwen4_exp as modeling |
| for name in ['causal_conv1d_fn','causal_conv1d_update','torch_chunk_gated_delta_rule','torch_recurrent_gated_delta_rule']: |
| f=getattr(modeling,name);setattr(modeling,name,getattr(f,'__wrapped__',f)) |
| torch.set_num_threads(8) |
| base=Path(__file__).resolve().parent |
| t=AutoTokenizer.from_pretrained(base,local_files_only=True);c=AutoConfig.from_pretrained(base,local_files_only=True);c.experts_impl='eager';c.output_router_logits=False |
| m=AutoModelForCausalLM.from_pretrained(base,config=c,dtype=torch.bfloat16,attn_implementation='eager',device_map=a.device,local_files_only=True).eval() |
| s=a.prompt if a.raw else t.apply_chat_template([{'role':'user','content':a.prompt}],tokenize=False,add_generation_prompt=True,enable_thinking=False);inputs=t(s,return_tensors='pt').to(a.device) |
| with torch.inference_mode(): |
| ids=m.generate(**inputs,max_new_tokens=a.max_new_tokens,do_sample=False,use_cache=True,logits_to_keep=1,pad_token_id=t.pad_token_id)[0,inputs.input_ids.shape[1]:] |
| print(t.decode(ids,skip_special_tokens=True)) |
|
|
