521 lines
20 KiB
Python
521 lines
20 KiB
Python
"""Unit tests for srt/sampling/penaltylib/ — no server, no model loading."""
|
|
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
|
|
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
from sglang.srt.sampling.penaltylib.frequency_penalty import (
|
|
BatchedFrequencyPenalizer,
|
|
)
|
|
from sglang.srt.sampling.penaltylib.min_new_tokens import (
|
|
BatchedMinNewTokensPenalizer,
|
|
)
|
|
from sglang.srt.sampling.penaltylib.orchestrator import (
|
|
BatchedPenalizerOrchestrator,
|
|
)
|
|
from sglang.srt.sampling.penaltylib.presence_penalty import (
|
|
BatchedPresencePenalizer,
|
|
)
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
VOCAB_SIZE = 32
|
|
DEVICE = "cpu"
|
|
|
|
|
|
# Helpers: mock Req and ScheduleBatch
|
|
def _make_req(freq=0.0, presence=0.0, min_tokens=0, stop_ids=None, eos_id=2):
|
|
"""Create a mock request with sampling params."""
|
|
req = MagicMock()
|
|
req.sampling_params.frequency_penalty = freq
|
|
req.sampling_params.presence_penalty = presence
|
|
req.sampling_params.min_new_tokens = min_tokens
|
|
req.sampling_params.stop_token_ids = stop_ids
|
|
req.tokenizer.additional_stop_token_ids = None
|
|
req.tokenizer.eos_token_id = eos_id
|
|
return req
|
|
|
|
|
|
def _make_batch(reqs):
|
|
"""Create a mock ScheduleBatch.
|
|
Note: orchestrator accesses batch.reqs as an attribute (not a method call)."""
|
|
batch = MagicMock()
|
|
batch.reqs = reqs
|
|
batch.device = DEVICE
|
|
return batch
|
|
|
|
|
|
# BatchedPenalizerOrchestrator
|
|
class TestBatchedPenalizerOrchestrator(CustomTestCase):
|
|
|
|
def test_init_detects_required_penalizers(self):
|
|
"""Test that orchestrator marks is_required=True when any request has nonzero penalty."""
|
|
reqs = [_make_req(freq=1.0)]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
self.assertTrue(orch.is_required)
|
|
|
|
def test_init_not_required_when_no_penalties(self):
|
|
"""Test that orchestrator marks is_required=False when all penalties are zero."""
|
|
reqs = [_make_req()] # all defaults (0.0)
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
self.assertFalse(orch.is_required)
|
|
|
|
def test_batch_property_via_weakref(self):
|
|
"""Test that batch property returns the original batch via weakref."""
|
|
reqs = [_make_req()]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(VOCAB_SIZE, batch, set())
|
|
self.assertIs(orch.batch, batch)
|
|
|
|
def test_batch_setter_none(self):
|
|
"""Test that setting batch to None breaks the weakref cleanly."""
|
|
reqs = [_make_req()]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(VOCAB_SIZE, batch, set())
|
|
orch.batch = None
|
|
self.assertIsNone(orch.batch)
|
|
|
|
def test_batch_setter_new_batch(self):
|
|
"""Test that batch can be reassigned to a different ScheduleBatch."""
|
|
reqs = [_make_req()]
|
|
batch1 = _make_batch(reqs)
|
|
batch2 = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(VOCAB_SIZE, batch1, set())
|
|
orch.batch = batch2
|
|
self.assertIs(orch.batch, batch2)
|
|
|
|
def test_context_manager_releases(self):
|
|
"""Test that exiting the context manager releases all penalizers."""
|
|
reqs = [_make_req(freq=1.0)]
|
|
batch = _make_batch(reqs)
|
|
with BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
) as orch:
|
|
self.assertTrue(orch.is_required)
|
|
self.assertFalse(orch.is_required)
|
|
self.assertEqual(len(orch.penalizers), 0)
|
|
|
|
def test_filter_empty_indices_releases(self):
|
|
"""Test that filtering with no indices left fully releases the orchestrator."""
|
|
reqs = [_make_req(freq=1.0)]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
orch.filter(torch.tensor([], dtype=torch.long))
|
|
self.assertFalse(orch.is_required)
|
|
|
|
def test_filter_not_required_is_noop(self):
|
|
"""Test that filter on a not-required orchestrator does nothing."""
|
|
reqs = [_make_req()]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
self.assertFalse(orch.is_required)
|
|
orch.filter(torch.tensor([0])) # should not raise
|
|
|
|
def test_merge_both_not_required_is_noop(self):
|
|
"""Test that merging two not-required orchestrators stays not-required."""
|
|
reqs = [_make_req()]
|
|
batch = _make_batch(reqs)
|
|
orch1 = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
orch2 = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
orch1.merge(orch2) # should not raise
|
|
self.assertFalse(orch1.is_required)
|
|
|
|
|
|
# BatchedFrequencyPenalizer
|
|
class TestBatchedFrequencyPenalizer(CustomTestCase):
|
|
|
|
def _setup(self, freq_values):
|
|
reqs = [_make_req(freq=f) for f in freq_values]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
pen = orch.penalizers[BatchedFrequencyPenalizer]
|
|
return orch, pen
|
|
|
|
def test_is_required_with_nonzero_penalty(self):
|
|
"""Test that nonzero frequency_penalty makes the penalizer required."""
|
|
_, pen = self._setup([1.5])
|
|
self.assertTrue(pen.is_required())
|
|
|
|
def test_is_not_required_with_zero_penalty(self):
|
|
"""Test that zero frequency_penalty makes the penalizer not required."""
|
|
_, pen = self._setup([0.0])
|
|
self.assertFalse(pen.is_required())
|
|
|
|
def test_cumulate_and_apply(self):
|
|
"""Test that cumulating a token applies frequency penalty to its logit."""
|
|
orch, pen = self._setup([2.0])
|
|
output_ids = torch.tensor([5])
|
|
pen.cumulate_output_tokens(output_ids)
|
|
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
self.assertAlmostEqual(logits[0, 5].item(), -2.0, places=5)
|
|
# Other tokens unaffected
|
|
self.assertAlmostEqual(logits[0, 0].item(), 0.0, places=5)
|
|
|
|
def test_cumulate_twice_doubles_penalty(self):
|
|
"""Test that frequency penalty scales linearly with occurrence count."""
|
|
orch, pen = self._setup([1.0])
|
|
pen.cumulate_output_tokens(torch.tensor([3]))
|
|
pen.cumulate_output_tokens(torch.tensor([3]))
|
|
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
self.assertAlmostEqual(logits[0, 3].item(), -2.0, places=5)
|
|
|
|
def test_filter_keeps_subset(self):
|
|
"""Test that filter retains only the selected batch indices."""
|
|
orch, pen = self._setup([1.0, 2.0])
|
|
keep = torch.tensor([1])
|
|
pen.filter(keep)
|
|
self.assertEqual(pen.frequency_penalties.shape[0], 1)
|
|
self.assertAlmostEqual(pen.frequency_penalties[0, 0].item(), 2.0, places=5)
|
|
|
|
def test_merge_concatenates(self):
|
|
"""Test that merge concatenates penalty tensors from two penalizers."""
|
|
_, pen1 = self._setup([1.0])
|
|
_, pen2 = self._setup([2.0])
|
|
pen1.merge(pen2)
|
|
self.assertEqual(pen1.frequency_penalties.shape[0], 2)
|
|
|
|
def test_teardown_cleans_attributes(self):
|
|
"""Test that teardown deletes internal tensors and resets prepared state."""
|
|
_, pen = self._setup([1.0])
|
|
pen.teardown()
|
|
self.assertFalse(hasattr(pen, "frequency_penalties"))
|
|
self.assertFalse(hasattr(pen, "cumulated_frequency_penalties"))
|
|
self.assertFalse(pen.is_prepared())
|
|
|
|
def test_cumulate_when_not_prepared_is_noop(self):
|
|
"""Test that cumulate before prepare does not crash."""
|
|
_, pen = self._setup([0.0])
|
|
# pen is not prepared (is_required=False)
|
|
pen.cumulate_output_tokens(torch.tensor([1])) # should not raise
|
|
|
|
def test_apply_when_not_prepared_is_noop(self):
|
|
"""Test that apply on an unprepared penalizer leaves logits unchanged."""
|
|
_, pen = self._setup([0.0])
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
original = logits.clone()
|
|
pen.apply(logits)
|
|
self.assertTrue(torch.equal(logits, original))
|
|
|
|
|
|
# BatchedPresencePenalizer
|
|
class TestBatchedPresencePenalizer(CustomTestCase):
|
|
|
|
def _setup(self, presence_values):
|
|
reqs = [_make_req(presence=p) for p in presence_values]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedPresencePenalizer}
|
|
)
|
|
pen = orch.penalizers[BatchedPresencePenalizer]
|
|
return orch, pen
|
|
|
|
def test_is_required_with_nonzero_penalty(self):
|
|
"""Test that nonzero presence_penalty makes the penalizer required."""
|
|
_, pen = self._setup([0.5])
|
|
self.assertTrue(pen.is_required())
|
|
|
|
def test_presence_penalty_does_not_scale(self):
|
|
"""Test that presence penalty is flat (same value regardless of count)."""
|
|
orch, pen = self._setup([1.0])
|
|
pen.cumulate_output_tokens(torch.tensor([7]))
|
|
pen.cumulate_output_tokens(torch.tensor([7])) # same token again
|
|
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
# scatter_ overwrites (not adds), so penalty should be 1.0, not 2.0
|
|
self.assertAlmostEqual(logits[0, 7].item(), -1.0, places=5)
|
|
|
|
def test_filter_keeps_subset(self):
|
|
"""Test that filter retains the first request's presence penalty."""
|
|
orch, pen = self._setup([1.0, 2.0])
|
|
keep = torch.tensor([0])
|
|
pen.filter(keep)
|
|
self.assertEqual(pen.presence_penalties.shape[0], 1)
|
|
self.assertAlmostEqual(pen.presence_penalties[0, 0].item(), 1.0, places=5)
|
|
|
|
def test_merge_concatenates(self):
|
|
"""Test that merge concatenates presence penalty tensors."""
|
|
_, pen1 = self._setup([1.0])
|
|
_, pen2 = self._setup([2.0])
|
|
pen1.merge(pen2)
|
|
self.assertEqual(pen1.presence_penalties.shape[0], 2)
|
|
|
|
def test_teardown_cleans_attributes(self):
|
|
"""Test that teardown removes the presence_penalties tensor."""
|
|
_, pen = self._setup([1.0])
|
|
pen.teardown()
|
|
self.assertFalse(hasattr(pen, "presence_penalties"))
|
|
|
|
|
|
# BatchedMinNewTokensPenalizer
|
|
class TestBatchedMinNewTokensPenalizer(CustomTestCase):
|
|
|
|
def _setup(self, configs):
|
|
"""configs: list of (min_tokens, stop_ids, eos_id)."""
|
|
reqs = [_make_req(min_tokens=c[0], stop_ids=c[1], eos_id=c[2]) for c in configs]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedMinNewTokensPenalizer}
|
|
)
|
|
pen = orch.penalizers[BatchedMinNewTokensPenalizer]
|
|
return orch, pen
|
|
|
|
def test_is_required_with_positive_min_tokens(self):
|
|
"""Test that positive min_new_tokens makes the penalizer required."""
|
|
_, pen = self._setup([(5, None, 2)])
|
|
self.assertTrue(pen.is_required())
|
|
|
|
def test_is_not_required_with_zero_min_tokens(self):
|
|
"""Test that min_new_tokens=0 makes the penalizer not required."""
|
|
_, pen = self._setup([(0, None, 2)])
|
|
self.assertFalse(pen.is_required())
|
|
|
|
def test_blocks_eos_before_min_tokens(self):
|
|
"""Test that EOS token is blocked before min_new_tokens is reached."""
|
|
orch, pen = self._setup([(3, None, 2)])
|
|
# Before any output: len=0 < min=3 → block EOS (token 2)
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
self.assertTrue(torch.isinf(logits[0, 2]) and logits[0, 2] < 0)
|
|
# Non-stop tokens should be fine
|
|
self.assertEqual(logits[0, 0].item(), 0.0)
|
|
|
|
def test_allows_eos_after_min_tokens(self):
|
|
"""Test that EOS is allowed after generating min_new_tokens."""
|
|
orch, pen = self._setup([(2, None, 2)])
|
|
# Generate 2 tokens
|
|
pen.cumulate_output_tokens(torch.tensor([10]))
|
|
pen.cumulate_output_tokens(torch.tensor([11]))
|
|
# Now len=2 >= min=2 → EOS should NOT be blocked
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
self.assertEqual(logits[0, 2].item(), 0.0)
|
|
|
|
def test_blocks_custom_stop_tokens(self):
|
|
"""Test that custom stop_token_ids are also blocked before min_new_tokens."""
|
|
orch, pen = self._setup([(3, {5, 10}, 2)])
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
# EOS (2), stop token 5, stop token 10 should all be blocked
|
|
self.assertTrue(torch.isinf(logits[0, 2]) and logits[0, 2] < 0)
|
|
self.assertTrue(torch.isinf(logits[0, 5]) and logits[0, 5] < 0)
|
|
self.assertTrue(torch.isinf(logits[0, 10]) and logits[0, 10] < 0)
|
|
|
|
def test_blocks_additional_stop_tokens(self):
|
|
"""Test that tokenizer's additional_stop_token_ids are also blocked."""
|
|
req = _make_req(min_tokens=3, stop_ids=None, eos_id=2)
|
|
req.tokenizer.additional_stop_token_ids = {7, 8}
|
|
batch = _make_batch([req])
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedMinNewTokensPenalizer}
|
|
)
|
|
pen = orch.penalizers[BatchedMinNewTokensPenalizer]
|
|
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
pen.apply(logits)
|
|
# EOS (2) + additional stops (7, 8) should all be blocked
|
|
for tok in [2, 7, 8]:
|
|
self.assertTrue(
|
|
torch.isinf(logits[0, tok]) and logits[0, tok] < 0,
|
|
f"token {tok} should be blocked before min_new_tokens",
|
|
)
|
|
# Non-stop tokens should be fine
|
|
self.assertEqual(logits[0, 0].item(), 0.0)
|
|
|
|
def test_filter_keeps_subset(self):
|
|
"""Test that filter keeps the second request (min_tokens=5) and drops the first."""
|
|
orch, pen = self._setup([(3, None, 2), (5, None, 2)])
|
|
keep = torch.tensor([1])
|
|
pen.filter(keep)
|
|
self.assertEqual(pen.min_new_tokens.shape[0], 1)
|
|
self.assertEqual(pen.min_new_tokens[0, 0].item(), 5)
|
|
|
|
def test_merge_concatenates(self):
|
|
"""Test that merge combines min_new_tokens tensors from two penalizers."""
|
|
_, pen1 = self._setup([(3, None, 2)])
|
|
_, pen2 = self._setup([(5, None, 2)])
|
|
pen1.merge(pen2)
|
|
self.assertEqual(pen1.min_new_tokens.shape[0], 2)
|
|
|
|
def test_teardown_cleans_attributes(self):
|
|
"""Test that teardown removes min_new_tokens, stop_token_penalties, and len_output_tokens."""
|
|
_, pen = self._setup([(3, None, 2)])
|
|
pen.teardown()
|
|
self.assertFalse(hasattr(pen, "min_new_tokens"))
|
|
self.assertFalse(hasattr(pen, "stop_token_penalties"))
|
|
self.assertFalse(hasattr(pen, "len_output_tokens"))
|
|
|
|
|
|
# _BatchedPenalizer base class edge cases
|
|
class TestBatchedPenalizerBase(CustomTestCase):
|
|
|
|
def test_filter_when_not_prepared_is_noop(self):
|
|
"""Test that filter on an unprepared penalizer does not crash."""
|
|
reqs = [_make_req()]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
pen = orch.penalizers[BatchedFrequencyPenalizer]
|
|
# pen is not prepared (frequency_penalty=0 → not required)
|
|
pen.filter(torch.tensor([0])) # should not raise
|
|
|
|
def test_merge_prepares_both_if_needed(self):
|
|
"""Test that merge prepares unprepared side before concatenating."""
|
|
reqs_a = [_make_req(freq=0.0)] # not required
|
|
reqs_b = [_make_req(freq=1.0)] # required
|
|
batch_a = _make_batch(reqs_a)
|
|
batch_b = _make_batch(reqs_b)
|
|
orch_a = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch_a, {BatchedFrequencyPenalizer}
|
|
)
|
|
orch_b = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch_b, {BatchedFrequencyPenalizer}
|
|
)
|
|
pen_a = orch_a.penalizers[BatchedFrequencyPenalizer]
|
|
pen_b = orch_b.penalizers[BatchedFrequencyPenalizer]
|
|
self.assertFalse(pen_a.is_prepared())
|
|
self.assertTrue(pen_b.is_prepared())
|
|
# Merge should prepare pen_a first
|
|
pen_a.merge(pen_b)
|
|
self.assertTrue(pen_a.is_prepared())
|
|
self.assertEqual(pen_a.frequency_penalties.shape[0], 2)
|
|
|
|
def test_merge_both_unprepared_is_noop(self):
|
|
"""Test that merging two unprepared penalizers keeps them unprepared."""
|
|
reqs = [_make_req()]
|
|
batch = _make_batch(reqs)
|
|
orch1 = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
orch2 = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
pen1 = orch1.penalizers[BatchedFrequencyPenalizer]
|
|
pen2 = orch2.penalizers[BatchedFrequencyPenalizer]
|
|
pen1.merge(pen2) # both not prepared → noop
|
|
self.assertFalse(pen1.is_prepared())
|
|
|
|
def test_prepare_is_idempotent(self):
|
|
"""Test that calling prepare() multiple times does not crash."""
|
|
reqs = [_make_req(freq=1.0)]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
pen = orch.penalizers[BatchedFrequencyPenalizer]
|
|
self.assertTrue(pen.is_prepared())
|
|
# Calling prepare again should not crash or reinitialize
|
|
pen.prepare()
|
|
self.assertTrue(pen.is_prepared())
|
|
|
|
|
|
# Orchestrator with multiple penalizer types
|
|
class TestOrchestratorMultiplePenalizers(CustomTestCase):
|
|
|
|
def test_all_three_penalizers(self):
|
|
"""Test orchestrator managing frequency, presence, and min_new_tokens together."""
|
|
reqs = [_make_req(freq=1.0, presence=0.5, min_tokens=2, eos_id=2)]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE,
|
|
batch,
|
|
{
|
|
BatchedFrequencyPenalizer,
|
|
BatchedPresencePenalizer,
|
|
BatchedMinNewTokensPenalizer,
|
|
},
|
|
)
|
|
self.assertTrue(orch.is_required)
|
|
|
|
# Cumulate one token
|
|
output_ids = torch.tensor([5])
|
|
orch.cumulate_output_tokens(output_ids)
|
|
|
|
# Apply all penalties
|
|
logits = torch.zeros(1, VOCAB_SIZE)
|
|
orch.apply(logits)
|
|
|
|
# Token 5: freq_penalty=1.0 (cumulated once) + pres_penalty=0.5
|
|
self.assertAlmostEqual(logits[0, 5].item(), -1.5, places=4)
|
|
# EOS (token 2): blocked by min_new_tokens (len=1 < min=2)
|
|
self.assertTrue(torch.isinf(logits[0, 2]) and logits[0, 2] < 0)
|
|
|
|
def test_filter_with_penalizer_no_longer_required(self):
|
|
"""Test that penalizer is torn down when no longer required after filter."""
|
|
reqs = [_make_req(freq=0.0), _make_req(freq=1.0)]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
self.assertTrue(orch.is_required)
|
|
|
|
# Keep only the request with freq=0 (index 0)
|
|
batch.reqs = [reqs[0]]
|
|
orch.filter(torch.tensor([0]))
|
|
|
|
pen = orch.penalizers[BatchedFrequencyPenalizer]
|
|
# After filter, only req with freq=0 remains → penalizer not required
|
|
self.assertFalse(pen.is_required())
|
|
|
|
def test_filter_keeps_required_penalizer(self):
|
|
"""Test that filter keeps penalizer active when still required."""
|
|
reqs = [_make_req(freq=1.0), _make_req(freq=2.0)]
|
|
batch = _make_batch(reqs)
|
|
orch = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch, {BatchedFrequencyPenalizer}
|
|
)
|
|
self.assertTrue(orch.is_required)
|
|
|
|
batch.reqs = [reqs[1]]
|
|
orch.filter(torch.tensor([1]))
|
|
self.assertTrue(orch.is_required)
|
|
|
|
def test_merge_one_required(self):
|
|
"""Test that merge marks orchestrator as required when one side is."""
|
|
reqs_a = [_make_req(freq=0.0)]
|
|
reqs_b = [_make_req(freq=1.0)]
|
|
batch_a = _make_batch(reqs_a)
|
|
batch_b = _make_batch(reqs_b)
|
|
orch_a = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch_a, {BatchedFrequencyPenalizer}
|
|
)
|
|
orch_b = BatchedPenalizerOrchestrator(
|
|
VOCAB_SIZE, batch_b, {BatchedFrequencyPenalizer}
|
|
)
|
|
self.assertFalse(orch_a.is_required)
|
|
self.assertTrue(orch_b.is_required)
|
|
|
|
orch_a.merge(orch_b)
|
|
self.assertTrue(orch_a.is_required)
|
|
pen = orch_a.penalizers[BatchedFrequencyPenalizer]
|
|
self.assertEqual(pen.frequency_penalties.shape[0], 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|