Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4
5# 1. load model
6device = "cuda" if torch.cuda.is_available() else "cpu"
7repo_id = "SakanaAI/TinySwallow-1.5B-Instruct"
8model = AutoModelForCausalLM.from_pretrained(repo_id)
9tokenizer = AutoTokenizer.from_pretrained(repo_id)
10model.to(device)
11
12# 2. prepare inputs
13text = "知識蒸留について簡単に教えてください。"
14messages = [{"role": "user", "content": text}]
15input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
16
17# 3. generate
18output_ids = model.generate(
19 input_ids.to(device),
20 max_new_tokens=1024,
21)
22output_ids = output_ids[:, input_ids.shape[1] :]
23generated_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
24print(generated_text)1@misc{sakana2025taid,
2 title = {TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models},
3 author. = {Makoto Shing and Kou Misaki and Han Bao and Sho Yokoi and Takuya Akiba},
4 year = {2025},
5 eprint = {2501.16937},
6 archivePrefix = {arXiv},
7 primaryClass = {cs.LG},
8 url = {https://arxiv.org/abs/2501.16937}
9}