Update README.md
Browse files
README.md
CHANGED
@@ -41,14 +41,14 @@ model.to(device)
|
|
41 |
# 2. prepare inputs
|
42 |
text = "知識蒸留について簡単に教えてください。"
|
43 |
messages = [{"role": "user", "content": text}]
|
44 |
-
|
45 |
|
46 |
# 3. generate
|
47 |
output_ids = model.generate(
|
48 |
inputs.to(device),
|
49 |
max_new_tokens=1024,
|
50 |
)
|
51 |
-
output_ids = output_ids[:,
|
52 |
generated_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
|
53 |
print(generated_text)
|
54 |
```
|
|
|
41 |
# 2. prepare inputs
|
42 |
text = "知識蒸留について簡単に教えてください。"
|
43 |
messages = [{"role": "user", "content": text}]
|
44 |
+
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
|
45 |
|
46 |
# 3. generate
|
47 |
output_ids = model.generate(
|
48 |
inputs.to(device),
|
49 |
max_new_tokens=1024,
|
50 |
)
|
51 |
+
output_ids = output_ids[:, input_ids.shape[1] :]
|
52 |
generated_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
|
53 |
print(generated_text)
|
54 |
```
|