mirror of
https://github.com/NVIDIA-NeMo/DataDesigner
synced 2026-05-24 09:48:29 +00:00
65 lines
2.8 KiB
Python
65 lines
2.8 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
|
|
|
|
|
|
def test_token_usage_stats():
|
|
token_usage_stats = TokenUsageStats()
|
|
assert token_usage_stats.prompt_tokens == 0
|
|
assert token_usage_stats.completion_tokens == 0
|
|
assert token_usage_stats.total_tokens == 0
|
|
assert token_usage_stats.has_usage is False
|
|
|
|
token_usage_stats.extend(prompt_tokens=10, completion_tokens=20)
|
|
assert token_usage_stats.prompt_tokens == 10
|
|
assert token_usage_stats.completion_tokens == 20
|
|
assert token_usage_stats.total_tokens == 30
|
|
assert token_usage_stats.has_usage is True
|
|
|
|
|
|
def test_request_usage_stats():
|
|
request_usage_stats = RequestUsageStats()
|
|
assert request_usage_stats.successful_requests == 0
|
|
assert request_usage_stats.failed_requests == 0
|
|
assert request_usage_stats.total_requests == 0
|
|
assert request_usage_stats.has_usage is False
|
|
|
|
request_usage_stats.extend(successful_requests=10, failed_requests=20)
|
|
assert request_usage_stats.successful_requests == 10
|
|
assert request_usage_stats.failed_requests == 20
|
|
assert request_usage_stats.total_requests == 30
|
|
assert request_usage_stats.has_usage is True
|
|
|
|
|
|
def test_model_usage_stats():
|
|
model_usage_stats = ModelUsageStats()
|
|
assert model_usage_stats.token_usage.prompt_tokens == 0
|
|
assert model_usage_stats.token_usage.completion_tokens == 0
|
|
assert model_usage_stats.request_usage.successful_requests == 0
|
|
assert model_usage_stats.request_usage.failed_requests == 0
|
|
assert model_usage_stats.has_usage is False
|
|
|
|
assert model_usage_stats.get_usage_stats(total_time_elapsed=10) == {
|
|
"token_usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
|
"request_usage": {"successful_requests": 0, "failed_requests": 0, "total_requests": 0},
|
|
"tokens_per_second": 0,
|
|
"requests_per_minute": 0,
|
|
}
|
|
|
|
model_usage_stats.extend(
|
|
token_usage=TokenUsageStats(prompt_tokens=10, completion_tokens=20),
|
|
request_usage=RequestUsageStats(successful_requests=2, failed_requests=1),
|
|
)
|
|
assert model_usage_stats.token_usage.prompt_tokens == 10
|
|
assert model_usage_stats.token_usage.completion_tokens == 20
|
|
assert model_usage_stats.request_usage.successful_requests == 2
|
|
assert model_usage_stats.request_usage.failed_requests == 1
|
|
assert model_usage_stats.has_usage is True
|
|
|
|
assert model_usage_stats.get_usage_stats(total_time_elapsed=2) == {
|
|
"token_usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
"request_usage": {"successful_requests": 2, "failed_requests": 1, "total_requests": 3},
|
|
"tokens_per_second": 15,
|
|
"requests_per_minute": 90,
|
|
}
|