gpted / expand_test.py
mebubo's picture
Better names for Expander etc
4537742
from dataclasses import dataclass
from expand import Series, BatchExpander, Expansion, Batch, TokenCandidates, BatchCandidates, CompletedSequence, CompletedBatch, expand
possible_sequences = [
[1, 21, 31, 41],
[1, 21, 31, 42],
[1, 21, 32, 41, 51],
[1, 22, 33, 41],
[1, 22, 34, 41],
]
def expand_series(series: Series) -> list[Expansion]:
all_tokens = series.get_all_tokens()
l = len(all_tokens)
items = [s[l] for s in possible_sequences if s[:l] == all_tokens and len(s) > l]
candidates = [Expansion(token=l, cost=-1.0) for l in dict.fromkeys(items)]
return candidates
class PredefinedSequenceExpander(BatchExpander):
def expand(self, batch: Batch) -> BatchCandidates:
result = []
for s in batch.items:
expansions = expand_series(s)
result.append(TokenCandidates(series=s, expansions=expansions))
return BatchCandidates(items=result)
expander = PredefinedSequenceExpander()
def test_expander_zero_budget():
s = Series(id=0, tokens=[1], budget=0.0)
expanded = expander.expand(Batch(items=[s]))
expected = BatchCandidates(
items=[TokenCandidates(series=s, expansions=[
Expansion(token=21, cost=-1.0),
Expansion(token=22, cost=-1.0),
])]
)
assert expected == expanded
def test_expander_budget_one():
s = Series(id=0, tokens=[1], budget=1.0)
expanded = expander.expand(Batch(items=[s]))
expected = BatchCandidates(
items=[TokenCandidates(series=s, expansions=[
Expansion(token=21, cost=-1.0),
Expansion(token=22, cost=-1.0),
])]
)
assert expected == expanded
def test_expander_budget_two():
s = Series(id=0, tokens=[1], budget=2.0)
expanded = expander.expand(Batch(items=[s]))
expected = BatchCandidates(
items=[TokenCandidates(series=s, expansions=[
Expansion(token=21, cost=-1.0),
Expansion(token=22, cost=-1.0),
])]
)
assert expected == expanded
def test_expander_budget_one_no_expansion():
s = Series(id=0, tokens=[1, 20], budget=1.0)
expanded = expander.expand(Batch(items=[s]))
expected = BatchCandidates(
items=[TokenCandidates(series=s, expansions=[])]
)
assert expected == expanded
def test_expander_budget_one_two_tokens():
s = Series(id=0, tokens=[1, 22], budget=1.0)
expanded = expander.expand(Batch(items=[s]))
expected = BatchCandidates(
items=[TokenCandidates(series=s, expansions=[
Expansion(token=33, cost=-1.0),
Expansion(token=34, cost=-1.0),
])]
)
assert expected == expanded
def test_expander_budget_one_two_tokens_two_series():
s1 = Series(id=0, tokens=[1, 21, 31], budget=1.0)
s2 = Series(id=1, tokens=[1, 22], budget=1.0)
expanded = expander.expand(Batch(items=[s1, s2]))
expected = BatchCandidates(
items=[
TokenCandidates(series=s1, expansions=[
Expansion(token=41, cost=-1.0),
Expansion(token=42, cost=-1.0),
]),
TokenCandidates(series=s2, expansions=[
Expansion(token=33, cost=-1.0),
Expansion(token=34, cost=-1.0),
])
]
)
assert expected == expanded
def test_expand_01():
batch = Batch(items=[
Series(id=0, tokens=[1, 21], budget=1.0),
Series(id=1, tokens=[1, 22], budget=1.0),
])
expanded = expand(batch, expander)
assert expanded == CompletedBatch(items=[
CompletedSequence(
series=Series(id=0, tokens=[1, 21], budget=1.0),
expansions=[
[Expansion(token=31, cost=-1.0)],
[Expansion(token=32, cost=-1.0)],
]
),
CompletedSequence(
series=Series(id=1, tokens=[1, 22], budget=1.0),
expansions=[
[Expansion(token=33, cost=-1.0)],
[Expansion(token=34, cost=-1.0)],
]
),
])
def test_expand_02():
batch = Batch(items=[
Series(id=0, tokens=[1, 21], budget=2.0),
Series(id=1, tokens=[1, 22], budget=1.0),
])
expanded = expand(batch, expander)
assert expanded == CompletedBatch(items=[
CompletedSequence(
series=Series(id=0, tokens=[1, 21], budget=2.0),
expansions=[
[Expansion(token=31, cost=-1.0), Expansion(token=41, cost=-1.0)],
[Expansion(token=31, cost=-1.0), Expansion(token=42, cost=-1.0)],
[Expansion(token=32, cost=-1.0), Expansion(token=41, cost=-1.0)],
]
),
CompletedSequence(
series=Series(id=1, tokens=[1, 22], budget=1.0),
expansions=[
[Expansion(token=33, cost=-1.0)],
[Expansion(token=34, cost=-1.0)],
]
),
])
def test_expand_03():
batch = Batch(items=[
Series(id=0, tokens=[1, 21], budget=3.0),
Series(id=1, tokens=[1, 22], budget=0.0),
])
expanded = expand(batch, expander)
assert expanded == CompletedBatch(items=[
CompletedSequence(
series=Series(id=0, tokens=[1, 21], budget=3.0),
expansions=[
[Expansion(token=31, cost=-1.0), Expansion(token=41, cost=-1.0)],
[Expansion(token=31, cost=-1.0), Expansion(token=42, cost=-1.0)],
[Expansion(token=32, cost=-1.0), Expansion(token=41, cost=-1.0), Expansion(token=51, cost=-1.0)],
]
),
CompletedSequence(
series=Series(id=1, tokens=[1, 22], budget=0.0),
expansions=[],
),
])