Better names for Expander etc

README.md

@@ -173,7 +173,7 @@ In my case, I stop when the budget is exhausted, and I also stop if the expansio
 
 Given the batch and the stopping criterion, we can call the expander:
 ```python
-expander = ExpanderOneBatchLLM(model, tokenizer)
+expander = LLMBatchExpander(model, tokenizer)
 expanded = expand(batch, expander, stopping_criterion)
 ```
 

completions.py

@@ -92,8 +92,7 @@ def check_text(input_text: str, model: PreTrainedModel, tokenizer: Tokenizer, de
 
     contexts = [word.context for _, word in low_prob_words]
 
-
-    expander = ExpanderOneBatchLLM(model, tokenizer)
+    expander = LLMBatchExpander(model, tokenizer)
 
     #%%
     series = []

expand.py

@@ -25,28 +25,28 @@ class Batch:
     items: list[Series]
 
 @dataclass
-class ExpansionOneResult:
+class TokenCandidates:
     series: Series
     expansions: list[Expansion]
 
 @dataclass
-class ExpansionOneResultBatch:
-    items: list[ExpansionOneResult]
+class BatchCandidates:
+    items: list[TokenCandidates]
 
 # A fundamental operation that we can implement both using an LLM and using a list of hardcoded sequences, for testing
-class ExpanderOneBatch(Protocol):
-    def expand(self, batch: Batch) -> ExpansionOneResultBatch: ...
+class BatchExpander(Protocol):
+    def expand(self, batch: Batch) -> BatchCandidates: ...
 
 @dataclass
-class ExpansionResult:
+class CompletedSequence:
     series: Series
     expansions: list[list[Expansion]]
 
 @dataclass
-class ExpansionResultBatch:
-    items: list[ExpansionResult]
+class CompletedBatch:
+    items: list[CompletedSequence]
 
-def compute_new_series(result: ExpansionOneResult, stopping_criterion: Callable[[Series, Expansion], bool]) -> tuple[list[Series], list[Series]]:
+def compute_new_series(result: TokenCandidates, stopping_criterion: Callable[[Series, Expansion], bool]) -> tuple[list[Series], list[Series]]:
     new_series_batch = []
     for expansion in result.expansions:
         if not stopping_criterion(result.series, expansion):
@@ -60,7 +60,7 @@ def compute_new_series(result: ExpansionOneResult, stopping_criterion: Callable[
     completed_series = [result.series] if len(new_series_batch) == 0 else []
     return new_series_batch, completed_series
 
-def compute_expansions(original_series: list[Series], expanded_series: list[Series]) -> ExpansionResultBatch:
+def compute_expansions(original_series: list[Series], expanded_series: list[Series]) -> CompletedBatch:
     # check that ids in original_series are unique
     assert len(original_series) == len({s.id for s in original_series})
     # group original series by id
@@ -73,15 +73,15 @@ def compute_expansions(original_series: list[Series], expanded_series: list[Seri
     results = []
     for id, s in original_series_by_id.items():
         expansions = expanded_series_by_id[id]
-        expansion_result = ExpansionResult(series=s, expansions=expansions)
+        expansion_result = CompletedSequence(series=s, expansions=expansions)
         results.append(expansion_result)
-    return ExpansionResultBatch(items=results)
+    return CompletedBatch(items=results)
 
 def default_completion_criterion(series: Series, expansion: Expansion) -> bool:
     return series.get_remaining_budget() + expansion.cost < 0
 
-# A compound operation that we can implement generically, relying on an ExpanderOneBatch
-def expand(batch: Batch, expander: ExpanderOneBatch, completion_criterion: Callable[[Series, Expansion], bool] = default_completion_criterion) -> ExpansionResultBatch:
+# A compound operation that we can implement generically, relying on a BatchExpander
+def expand(batch: Batch, expander: BatchExpander, completion_criterion: Callable[[Series, Expansion], bool] = default_completion_criterion) -> CompletedBatch:
     completed_series: list[Series] = []
     current_batch = batch
     while len(current_batch.items) > 0:

expand_llm.py

@@ -22,18 +22,18 @@ def prepare_inputs(contexts: list[list[int]], tokenizer: Tokenizer, device: torc
     return tokenizer(texts, return_tensors="pt", padding=True).to(device)
 
 @dataclass
-class ExpanderOneBatchLLM:
+class LLMBatchExpander(BatchExpander):
     model: PreTrainedModel
     tokenizer: Tokenizer
 
-    def expand(self, batch: Batch) -> ExpansionOneResultBatch:
+    def expand(self, batch: Batch) -> BatchCandidates:
         inputs = prepare_inputs([s.get_all_tokens() for s in batch.items], self.tokenizer, self.model.device)
         next_tokens = find_next_tokens(self.model, inputs, self.tokenizer)
         results = []
         for s, next_tokens in zip(batch.items, next_tokens):
             expansions = [Expansion(token=token, cost=cost) for token, cost in next_tokens]
-            results.append(ExpansionOneResult(series=s, expansions=expansions))
-        return ExpansionOneResultBatch(items=results)
+            results.append(TokenCandidates(series=s, expansions=expansions))
+        return BatchCandidates(items=results)
 
 def create_stopping_criterion_llm(tokenizer: Tokenizer) -> Callable[[Series, Expansion], bool]:
     def stopping_criterion(series: Series, expansion: Expansion) -> bool:

expand_test.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from expand import Series, ExpanderOneBatch, Expansion, Batch, ExpansionOneResult, ExpansionOneResultBatch, ExpansionResult, ExpansionResultBatch, expand
+from expand import Series, BatchExpander, Expansion, Batch, TokenCandidates, BatchCandidates, CompletedSequence, CompletedBatch, expand
 
 possible_sequences = [
     [1, 21, 31, 41],
@@ -16,21 +16,21 @@ def expand_series(series: Series) -> list[Expansion]:
     candidates = [Expansion(token=l, cost=-1.0) for l in dict.fromkeys(items)]
     return candidates
 
-class HardcodedExpanderOneBatch(ExpanderOneBatch):
-    def expand(self, batch: Batch) -> ExpansionOneResultBatch:
+class PredefinedSequenceExpander(BatchExpander):
+    def expand(self, batch: Batch) -> BatchCandidates:
         result = []
         for s in batch.items:
             expansions = expand_series(s)
-            result.append(ExpansionOneResult(series=s, expansions=expansions))
-        return ExpansionOneResultBatch(items=result)
+            result.append(TokenCandidates(series=s, expansions=expansions))
+        return BatchCandidates(items=result)
 
-expander = HardcodedExpanderOneBatch()
+expander = PredefinedSequenceExpander()
 
 def test_expander_zero_budget():
     s = Series(id=0, tokens=[1], budget=0.0)
     expanded = expander.expand(Batch(items=[s]))
-    expected = ExpansionOneResultBatch(
-        items=[ExpansionOneResult(series=s, expansions=[
+    expected = BatchCandidates(
+        items=[TokenCandidates(series=s, expansions=[
             Expansion(token=21, cost=-1.0),
             Expansion(token=22, cost=-1.0),
         ])]
@@ -40,8 +40,8 @@ def test_expander_zero_budget():
 def test_expander_budget_one():
     s = Series(id=0, tokens=[1], budget=1.0)
     expanded = expander.expand(Batch(items=[s]))
-    expected = ExpansionOneResultBatch(
-        items=[ExpansionOneResult(series=s, expansions=[
+    expected = BatchCandidates(
+        items=[TokenCandidates(series=s, expansions=[
             Expansion(token=21, cost=-1.0),
             Expansion(token=22, cost=-1.0),
         ])]
@@ -51,8 +51,8 @@ def test_expander_budget_one():
 def test_expander_budget_two():
     s = Series(id=0, tokens=[1], budget=2.0)
     expanded = expander.expand(Batch(items=[s]))
-    expected = ExpansionOneResultBatch(
-        items=[ExpansionOneResult(series=s, expansions=[
+    expected = BatchCandidates(
+        items=[TokenCandidates(series=s, expansions=[
             Expansion(token=21, cost=-1.0),
             Expansion(token=22, cost=-1.0),
         ])]
@@ -62,16 +62,16 @@ def test_expander_budget_two():
 def test_expander_budget_one_no_expansion():
     s = Series(id=0, tokens=[1, 20], budget=1.0)
     expanded = expander.expand(Batch(items=[s]))
-    expected = ExpansionOneResultBatch(
-        items=[ExpansionOneResult(series=s, expansions=[])]
+    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 = ExpansionOneResultBatch(
-        items=[ExpansionOneResult(series=s, expansions=[
+    expected = BatchCandidates(
+        items=[TokenCandidates(series=s, expansions=[
             Expansion(token=33, cost=-1.0),
             Expansion(token=34, cost=-1.0),
         ])]
@@ -82,13 +82,13 @@ 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 = ExpansionOneResultBatch(
+    expected = BatchCandidates(
         items=[
-            ExpansionOneResult(series=s1, expansions=[
+            TokenCandidates(series=s1, expansions=[
                 Expansion(token=41, cost=-1.0),
                 Expansion(token=42, cost=-1.0),
             ]),
-            ExpansionOneResult(series=s2, expansions=[
+            TokenCandidates(series=s2, expansions=[
                 Expansion(token=33, cost=-1.0),
                 Expansion(token=34, cost=-1.0),
             ])
@@ -102,15 +102,15 @@ def test_expand_01():
         Series(id=1, tokens=[1, 22], budget=1.0),
     ])
     expanded = expand(batch, expander)
-    assert expanded == ExpansionResultBatch(items=[
-        ExpansionResult(
+    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)],
             ]
         ),
-        ExpansionResult(
+        CompletedSequence(
             series=Series(id=1, tokens=[1, 22], budget=1.0),
             expansions=[
                 [Expansion(token=33, cost=-1.0)],
@@ -125,8 +125,8 @@ def test_expand_02():
         Series(id=1, tokens=[1, 22], budget=1.0),
     ])
     expanded = expand(batch, expander)
-    assert expanded == ExpansionResultBatch(items=[
-        ExpansionResult(
+    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)],
@@ -134,7 +134,7 @@ def test_expand_02():
                 [Expansion(token=32, cost=-1.0), Expansion(token=41, cost=-1.0)],
             ]
         ),
-        ExpansionResult(
+        CompletedSequence(
             series=Series(id=1, tokens=[1, 22], budget=1.0),
             expansions=[
                 [Expansion(token=33, cost=-1.0)],
@@ -149,8 +149,8 @@ def test_expand_03():
         Series(id=1, tokens=[1, 22], budget=0.0),
     ])
     expanded = expand(batch, expander)
-    assert expanded == ExpansionResultBatch(items=[
-        ExpansionResult(
+    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)],
@@ -158,7 +158,7 @@ def test_expand_03():
                 [Expansion(token=32, cost=-1.0), Expansion(token=41, cost=-1.0), Expansion(token=51, cost=-1.0)],
             ]
         ),
-        ExpansionResult(
+        CompletedSequence(
             series=Series(id=1, tokens=[1, 22], budget=0.0),
             expansions=[],
         ),

run.py

@@ -24,7 +24,7 @@ low_prob_words = [(i, word) for i, word in enumerate(words) if word.logprob < lo
 contexts = [word.context for _, word in low_prob_words]
 
 #%%
-expander = ExpanderOneBatchLLM(model, tokenizer)
+expander = LLMBatchExpander(model, tokenizer)
 
 #%%
 series = []
@@ -41,7 +41,7 @@ stopping_criterion = create_stopping_criterion_llm(tokenizer)
 expanded = expand(batch, expander, stopping_criterion)
 
 # %%
-def print_expansions(expansions: ExpansionResultBatch):
+def print_expansions(expansions: CompletedBatch):
     for result in expansions.items:
         for expansion in result.expansions:
             # convert tokens to string