import torchfrom peft import PeftModelfrom transformers import AutoModelForCausalLM, AutoTokenizer def build_forbidden_token_mask(tokenizer, vocab_size, forbidden_chars='eE', ascii_only=False, device=None, preserve_added_tokens=False): """Mask decoded token text while preserving special tokens such as EOS.""" special_ids = set(tokenizer.all_special_ids) if preserve_added_tokens: special_ids.update(tokenizer.added_tokens_decoder) texts = tokenizer.batch_decode([[i] for i in range(vocab_size)]) return torch.tensor([ i not in special_ids and ( any(c in forbidden_chars for c in text) or (ascii_only and not text.isascii()) ) for i, text in enumerate(texts) ], dtype=torch.bool, device=device) tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-1.7B')model = AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-1.7B', dtype='auto', device_map='auto')model = PeftModel.from_pretrained(model, 'k33wu/qwen3-1.7b-no-e-sft180') # the same mask used in testing; bad_words_ids blocks each of these tokensforbidden = build_forbidden_token_mask(tokenizer, model.config.vocab_size, 'eE', ascii_only=True)bad_words_ids = [[i] for i in forbidden.nonzero().flatten().tolist()] messages = [ {'role': 'system', 'content': "Answer in natural English without the letters 'e' or 'E'. " "Preserve correct spelling and answer the user's request. " "This restriction applies only to your final answer, not your thinking."}, {'role': 'user', 'content': 'Describe the ocean.'},]ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, enable_thinking=False, return_tensors='pt', return_dict=True).to(model.device)# temperature 0.7 over all tokens (no top-k/top-p), as in the eval belowout = model.generate(**ids, bad_words_ids=bad_words_ids, max_new_tokens=300, do_sample=True, temperature=0.7, top_k=0, top_p=1.0)print(tokenizer.decode(out[0, ids['input_ids'].shape[1]:], skip_special_tokens=True))