Views
No views yet
qwen3_0.6b+qwen3_4b_gsm8k_Fuser.1import torch
2from huggingface_hub import snapshot_download
3from script.playground.inference_example import load_rosetta_model, run_inference_example
4
5checkpoint_dir = snapshot_download(
6 repo_id="nics-efc/C2C_Fuser",
7 allow_patterns=["qwen3_0.6b+qwen2.5_0.5b_Fuser/*"],
8)
9
10model_config = {
11 "rosetta_config": {
12 "base_model": "Qwen/Qwen3-0.6B",
13 "teacher_model": "Qwen/Qwen2.5-0.5B-Instruct",
14 "checkpoints_dir": f"{checkpoint_dir}/qwen3_0.6b+qwen2.5_0.5b_Fuser/final",
15 }
16}
17
18rosetta_model, tokenizer = load_rosetta_model(model_config, eval_config={}, device=torch.device("cuda"))
19device = rosetta_model.device
20
21prompt = [{"role": "user", "content": "Say hello in one short sentence."}]
22input_text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True, enable_thinking=False)
23inputs = tokenizer(input_text, return_tensors="pt").to(device)
24
25instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(inputs['input_ids'].shape[1] - 1, 1).unsqueeze(0).to(device)
26label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(1, 1).unsqueeze(0).to(device)
27kv_cache_index = [instruction_index, label_index]
28
29with torch.no_grad():
30 sampling_params = {
31 'do_sample': False,
32 'max_new_tokens': 256
33 }
34 outputs = rosetta_model.generate(**inputs, kv_cache_index=kv_cache_index, **sampling_params)
35 output_text = tokenizer.decode(outputs[0, instruction_index.shape[1] + 1:], skip_special_tokens=True)
36 print(f"C2C output text: {output_text}")