Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions tests/models/hunyuan_v1_moe/test_modeling_hunyuan_v1_moe.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import unittest

import pytest
import torch
from parameterized import parameterized

from transformers import is_torch_available
Expand Down Expand Up @@ -110,10 +111,12 @@ def tearDown(self):
def test_model_generation(self):
# we will compele this when model file change over
# pass
EXPECTED_ANSWER = "\nOkay, I need to write a short summary about the benefits of regular exercise. Let me start by recalling what I know. First,"
EXPECTED_ANSWER = "\nOkay, I need to write a"
prompt = "Write a short summary of the benefits of regular exercise"
tokenizer = AutoTokenizer.from_pretrained("tencent/Hunyuan-A13B-Instruct")
model = AutoModelForCausalLM.from_pretrained("tencent/Hunyuan-A13B-Instruct", device_map="auto")
model = AutoModelForCausalLM.from_pretrained(
"tencent/Hunyuan-A13B-Instruct", device_map="auto", dtype=torch.bfloat16
)
messages = [
{"role": "user", "content": prompt},
]
Expand All @@ -123,7 +126,7 @@ def test_model_generation(self):
add_generation_prompt=True,
return_tensors="pt",
)
generated_ids = model.generate(tokenized_chat.to(model.device), max_new_tokens=30, top_k=1)
generated_ids = model.generate(tokenized_chat.to(model.device), max_new_tokens=10, top_k=1)
text = tokenizer.decode(generated_ids[0])
output = text.split("<think>")[1]
self.assertEqual(EXPECTED_ANSWER, output)