Snapshot

expand.py

@@ -1,26 +1,33 @@
 from collections import defaultdict
-from dataclasses import dataclass
-from typing import Protocol
+from dataclasses import dataclass, field
+from typing import Protocol, Self
+
+@dataclass
+class Expansion:
+    token: int
+    cost: float
 
 @dataclass
 class Series:
     id: int
     tokens: list[int]
     budget: float
+    expansions: list[Expansion] = field(default_factory=list)
+
+    def get_all_tokens(self) -> list[int]:
+        return self.tokens + [e.token for e in self.expansions]
+
+    def get_remaining_budget(self) -> float:
+        return self.budget + sum(e.cost for e in self.expansions)
 
 @dataclass
 class Batch:
     items: list[Series]
 
-@dataclass
-class ExpansionOne:
-    token: int
-    cost: float
-
 @dataclass
 class ExpansionOneResult:
     series: Series
-    expansions: list[ExpansionOne]
+    expansions: list[Expansion]
 
 @dataclass
 class ExpansionOneResultBatch:
@@ -33,7 +40,7 @@ class ExpanderOneBatch(Protocol):
 @dataclass
 class ExpansionResult:
     series: Series
-    expansions: list[list[int]]
+    expansions: list[list[Expansion]]
 
 @dataclass
 class ExpansionResultBatch:
@@ -42,7 +49,12 @@ class ExpansionResultBatch:
 def compute_new_series(result: ExpansionOneResult) -> list[Series]:
     results = []
     for expansion in result.expansions:
-        results.append(Series(id=result.series.id, tokens=result.series.tokens + [expansion.token], budget=result.series.budget + expansion.cost))
+        results.append(Series(
+            id=result.series.id,
+            tokens=result.series.tokens,
+            expansions=result.series.expansions + [expansion],
+            budget=result.series.budget
+        ))
     return results
 
 def compute_expansions(original_series: list[Series], expanded_series: list[Series]) -> ExpansionResultBatch:
@@ -51,16 +63,14 @@ def compute_expansions(original_series: list[Series], expanded_series: list[Seri
     # group original series by id
     original_series_by_id = {s.id: s for s in original_series}
     # group expanded series by id
-    expanded_series_by_id: dict[int, list[list[int]]] = defaultdict(list)
+    expanded_series_by_id: dict[int, list[list[Expansion]]] = defaultdict(list)
     for s in expanded_series:
-        expanded_series_by_id[s.id].append(s.tokens)
+        if len(s.expansions) != 0:
+            expanded_series_by_id[s.id].append(s.expansions)
     results = []
     for id, s in original_series_by_id.items():
         expansions = expanded_series_by_id[id]
-        # subtract the original series from each expansion
-        l = len(s.tokens)
-        trimmed_expansions = [e[l:] for e in expansions if len(e) > l]
-        expansion_result = ExpansionResult(series=s, expansions=trimmed_expansions)
+        expansion_result = ExpansionResult(series=s, expansions=expansions)
         results.append(expansion_result)
     return ExpansionResultBatch(items=results)
 

expand_llm.py

@@ -15,6 +15,6 @@ class ExpanderOneBatchLLM:
         next_tokens = find_next_tokens(self.model, inputs, self.tokenizer)
         results = []
         for s, next_tokens in zip(batch.items, next_tokens):
-            expansions = [ExpansionOne(token=token, cost=logprob) for token, logprob in next_tokens if logprob + s.budget >= 0]
+            expansions = [Expansion(token=token, cost=logprob) for token, logprob in next_tokens if logprob + s.get_remaining_budget() >= 0]
             results.append(ExpansionOneResult(series=s, expansions=expansions))
         return ExpansionOneResultBatch(items=results)

expand_test.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from expand import Series, ExpanderOneBatch, ExpansionOne, Batch, ExpansionOneResult, ExpansionOneResultBatch, ExpansionResult, ExpansionResultBatch, expand
+from expand import Series, ExpanderOneBatch, Expansion, Batch, ExpansionOneResult, ExpansionOneResultBatch, ExpansionResult, ExpansionResultBatch, expand
 
 possible_sequences = [
     [1, 21, 31, 41],
@@ -9,11 +9,12 @@ possible_sequences = [
     [1, 22, 34, 41],
 ]
 
-def expand_series(series: Series) -> list[ExpansionOne]:
-    l = len(series.tokens)
-    items = [s[l] for s in possible_sequences if s[:l] == series.tokens and len(s) > l]
-    candidates = [ExpansionOne(token=l, cost=-1.0) for l in dict.fromkeys(items)]
-    return [c for c in candidates if c.cost + series.budget >= 0]
+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 [c for c in candidates if c.cost + series.get_remaining_budget() >= 0]
 
 class HardcodedExpanderOneBatch(ExpanderOneBatch):
     def expand(self, batch: Batch) -> ExpansionOneResultBatch:
@@ -38,8 +39,8 @@ def test_expander_budget_one():
     expanded = expander.expand(Batch(items=[s]))
     expected = ExpansionOneResultBatch(
         items=[ExpansionOneResult(series=s, expansions=[
-            ExpansionOne(token=21, cost=-1.0),
-            ExpansionOne(token=22, cost=-1.0),
+            Expansion(token=21, cost=-1.0),
+            Expansion(token=22, cost=-1.0),
         ])]
     )
     assert expected == expanded
@@ -49,8 +50,8 @@ def test_expander_budget_two():
     expanded = expander.expand(Batch(items=[s]))
     expected = ExpansionOneResultBatch(
         items=[ExpansionOneResult(series=s, expansions=[
-            ExpansionOne(token=21, cost=-1.0),
-            ExpansionOne(token=22, cost=-1.0),
+            Expansion(token=21, cost=-1.0),
+            Expansion(token=22, cost=-1.0),
         ])]
     )
     assert expected == expanded
@@ -68,8 +69,8 @@ def test_expander_budget_one_two_tokens():
     expanded = expander.expand(Batch(items=[s]))
     expected = ExpansionOneResultBatch(
         items=[ExpansionOneResult(series=s, expansions=[
-            ExpansionOne(token=33, cost=-1.0),
-            ExpansionOne(token=34, cost=-1.0),
+            Expansion(token=33, cost=-1.0),
+            Expansion(token=34, cost=-1.0),
         ])]
     )
     assert expected == expanded
@@ -81,12 +82,12 @@ def test_expander_budget_one_two_tokens_two_series():
     expected = ExpansionOneResultBatch(
         items=[
             ExpansionOneResult(series=s1, expansions=[
-                ExpansionOne(token=41, cost=-1.0),
-                ExpansionOne(token=42, cost=-1.0),
+                Expansion(token=41, cost=-1.0),
+                Expansion(token=42, cost=-1.0),
             ]),
             ExpansionOneResult(series=s2, expansions=[
-                ExpansionOne(token=33, cost=-1.0),
-                ExpansionOne(token=34, cost=-1.0),
+                Expansion(token=33, cost=-1.0),
+                Expansion(token=34, cost=-1.0),
             ])
         ]
     )
@@ -102,15 +103,15 @@ def test_expand_01():
         ExpansionResult(
             series=Series(id=0, tokens=[1, 21], budget=1.0),
             expansions=[
-                [31],
-                [32],
+                [Expansion(token=31, cost=-1.0)],
+                [Expansion(token=32, cost=-1.0)],
             ]
         ),
         ExpansionResult(
             series=Series(id=1, tokens=[1, 22], budget=1.0),
             expansions=[
-                [33],
-                [34],
+                [Expansion(token=33, cost=-1.0)],
+                [Expansion(token=34, cost=-1.0)],
             ]
         ),
     ])
@@ -125,16 +126,16 @@ def test_expand_02():
         ExpansionResult(
             series=Series(id=0, tokens=[1, 21], budget=2.0),
             expansions=[
-                [31, 41],
-                [31, 42],
-                [32, 41],
+                [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)],
             ]
         ),
         ExpansionResult(
             series=Series(id=1, tokens=[1, 22], budget=1.0),
             expansions=[
-                [33],
-                [34],
+                [Expansion(token=33, cost=-1.0)],
+                [Expansion(token=34, cost=-1.0)],
             ]
         ),
     ])
@@ -149,9 +150,9 @@ def test_expand_03():
         ExpansionResult(
             series=Series(id=0, tokens=[1, 21], budget=3.0),
             expansions=[
-                [31, 41],
-                [31, 42],
-                [32, 41, 51],
+                [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)],
             ]
         ),
         ExpansionResult(

run.py

@@ -29,7 +29,7 @@ expander = ExpanderOneBatchLLM(model, tokenizer)
 #%%
 series = []
 for i, x in enumerate(contexts):
-    series.append(Series(id=i, tokens=x, budget=5.0))
+    series.append(Series(id=i, tokens=x, budget=7.0))
 
 #%%
 batch = Batch(items=series)
@@ -42,7 +42,8 @@ def print_expansions(expansions: ExpansionResultBatch):
     for result in expansions.items:
         for expansion in result.expansions:
             # convert tokens to string
-            s = tokenizer.decode(expansion)
+            tokens = [e.token for e in expansion]
+            s = tokenizer.decode(tokens)
             print(f"{result.series.id}: {expansion} {s}")
 
 print_expansions(expanded)