Views
No views yet

transformers, torch, and flash-attention installed. We used torch==2.10 and transformers==5.8.0.trust_remote_code=True to load the custom Orthrus architecture.1from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
2import torch
3
4MODEL_PATH = "chiennv/Orthrus-Qwen3-1.7B"
5
6# Load the model and tokenizer
7model = AutoModelForCausalLM.from_pretrained(
8 MODEL_PATH,
9 dtype=torch.bfloat16,
10 device_map="cuda",
11 attn_implementation="flash_attention_2", # options: sdpa | eager | flash_attention_4
12 trust_remote_code=True # Note: trust_remote_code=True is required
13).eval()
14tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
15
16prompt = "Write a program to count the frequency of each word in a paragraph."
17messages = [
18 {"role": "system", "content": ""},
19 {"role": "user", "content": prompt}
20]
21
22input_ids = tokenizer.apply_chat_template(
23 messages,
24 tokenize=True,
25 enable_thinking=False,
26 add_generation_prompt=True,
27 return_tensors="pt",
28).input_ids
29
30# Generate text natively utilizing parallel diffusion projection
31output_ids = model.generate(
32 input_ids=input_ids.to(model.device),
33 max_new_tokens=2048,
34 use_diffusion_mode=True,
35 streamer=TextStreamer(tokenizer, skip_prompt=True) # enable streaming
36)1@misc{vannguyen2026orthrusmemoryefficientparalleltoken,
2 title={Orthrus: Memory-Efficient Parallel Token Generation via Dual-View Diffusion},
3 author={Chien Van Nguyen and Chaitra Hegde and Van Cuong Pham and Ryan A. Rossi and Franck Dernoncourt and Thien Huu Nguyen},
4 year={2026},
5 eprint={2605.12825},
6 archivePrefix={arXiv},
7 primaryClass={cs.LG},
8 url={https://arxiv.org/abs/2605.12825},
9}