import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "aasim-m/DAFT-Qwen3-0.6B-checkpoint-4000"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id, torch_dtype=torch.bfloat16
).to("cuda").eval()
instruction = (
"Translate the following NVIDIA GPU assembly function into corresponding AMD GPU assembly. "
"Preserve the function's behavior. Return only the translated assembly, without explanations or Markdown fences."
)
source_assembly = "..."
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": instruction + "\n" + source_assembly}],
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
)
inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
remaining = 32768 - inputs["input_ids"].shape[1]
if remaining <= 0:
raise ValueError("Input leaves no output space within 32,768 tokens")
with torch.inference_mode():
output = model.generate(
**inputs,
max_new_tokens=min(4096, remaining),
do_sample=False,
use_cache=True,
pad_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))