TtT-3B

from modeling_qwen_TtT import Qwen2ForARDiffLM
from transformers import AutoTokenizer

# Load model and tokenizer
model = Qwen2ForARDiffLM.from_pretrained("Stephen-Lee/TtT-3B", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("Stephen-Lee/TtT-3B", trust_remote_code=True)

# Set up special tokens
if tokenizer.mask_token_id is None:
    tokenizer.mask_token_id = tokenizer.convert_tokens_to_ids("<|mask_token|>")

# Generate
prompt = (
    "<|im_start|>user\n"
    "<|begin_of_audio|><|audio_1234|>...<|end_of_audio|><|im_end|>\n"
    "<|im_start|>assistant\n"
)
output = generate(model, tokenizer, prompt, max_gen_len=2048)