Spaces:
Sleeping
Sleeping
| from unittest.mock import Mock | |
| import numpy as np | |
| from deep_learning.pipeline.impl.text_generation.generation import generate_with_stateful_model, generate_with_training_model | |
| from deep_learning.models.spec import GenerationContext | |
| def test_generate_with_training_model(): | |
| model = Mock() | |
| model.predict = Mock(return_value=np.zeros((1, 10, 100))) | |
| sample_results = [50, 60, 99] | |
| sample_fn = Mock(side_effect=[np.array([t]) for t in sample_results]) | |
| context = GenerationContext(end_of_text=99, max_length=10, sample_fn=sample_fn) | |
| result = generate_with_training_model(model, context, [10, 20]) | |
| assert result.token_ids == [10, 20, 50, 60] | |
| assert result.stop_reason == "<|endoftext|>" | |
| def test_generate_with_stateful_model(): | |
| model = Mock() | |
| model.predict = Mock( | |
| return_value=[np.zeros((1, 1, 100)), np.zeros((1, 16)), np.zeros((1, 16))] | |
| ) | |
| sample_results = [50, 60, 99] | |
| sample_fn = Mock(side_effect=[np.array([t]) for t in sample_results]) | |
| context = GenerationContext(end_of_text=99, max_length=10, sample_fn=sample_fn) | |
| initial_states = [np.zeros((1, 16)), np.zeros((1, 16))] | |
| result = generate_with_stateful_model(model, context, [10, 20], initial_states) | |
| assert result.token_ids == [10, 20, 50, 60] | |
| assert result.stop_reason == "<|endoftext|>" | |