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))