Working stopping on word boundary

expand.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Protocol, Self
+from typing import Callable, Protocol, Self
 
 @dataclass
 class Expansion:
@@ -46,16 +46,19 @@ class ExpansionResult:
 class ExpansionResultBatch:
     items: list[ExpansionResult]
 
-def compute_new_series(result: ExpansionOneResult) -> list[Series]:
-    results = []
+def compute_new_series(result: ExpansionOneResult, stopping_criterion: Callable[[Series, Expansion], bool]) -> tuple[list[Series], list[Series]]:
+    new_series_batch = []
     for expansion in result.expansions:
-        results.append(Series(
-            id=result.series.id,
-            tokens=result.series.tokens,
-            expansions=result.series.expansions + [expansion],
-            budget=result.series.budget
-        ))
-    return results
+        if not stopping_criterion(result.series, expansion):
+            new_series = Series(
+                id=result.series.id,
+                tokens=result.series.tokens,
+                expansions=result.series.expansions + [expansion],
+                budget=result.series.budget
+            )
+            new_series_batch.append(new_series)
+    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:
     # check that ids in original_series are unique
@@ -74,8 +77,11 @@ def compute_expansions(original_series: list[Series], expanded_series: list[Seri
         results.append(expansion_result)
     return ExpansionResultBatch(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) -> ExpansionResultBatch:
+def expand(batch: Batch, expander: ExpanderOneBatch, completion_criterion: Callable[[Series, Expansion], bool] = default_completion_criterion) -> ExpansionResultBatch:
     completed_series: list[Series] = []
     current_batch = batch
     while len(current_batch.items) > 0:
@@ -86,6 +92,8 @@ def expand(batch: Batch, expander: ExpanderOneBatch) -> ExpansionResultBatch:
             if len(item.expansions) == 0:
                 completed_series.append(item.series)
             else:
-                current_batch_items.extend(compute_new_series(item))
+                new_series, completed = compute_new_series(item, completion_criterion)
+                completed_series.extend(completed)
+                current_batch_items.extend(new_series)
         current_batch = Batch(items=current_batch_items)
     return compute_expansions(batch.items, completed_series)

expand_llm.py

@@ -15,6 +15,22 @@ class ExpanderOneBatchLLM:
         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 if cost + s.get_remaining_budget() >= 0]
+            expansions = [Expansion(token=token, cost=cost) for token, cost in next_tokens]
             results.append(ExpansionOneResult(series=s, expansions=expansions))
         return ExpansionOneResultBatch(items=results)
+
+def create_stopping_criterion_llm(tokenizer: Tokenizer) -> Callable[[Series, Expansion], bool]:
+    def stopping_criterion(series: Series, expansion: Expansion) -> bool:
+        d = default_completion_criterion(series, expansion)
+        if d:
+            return d
+        token_str = tokenizer.decode([expansion.token])
+        starts_with_space = token_str.startswith(" ")
+        print(f"-----{token_str}-----, {starts_with_space=}")
+        is_first_token = len(series.expansions) == 0
+        if is_first_token and not starts_with_space:
+            return True
+        if not is_first_token and starts_with_space:
+            return True
+        return False
+    return stopping_criterion

expand_test.py

@@ -14,7 +14,7 @@ def expand_series(series: Series) -> list[Expansion]:
     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]
+    return candidates
 
 class HardcodedExpanderOneBatch(ExpanderOneBatch):
     def expand(self, batch: Batch) -> ExpansionOneResultBatch:
@@ -30,7 +30,10 @@ 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=[])]
+        items=[ExpansionOneResult(series=s, expansions=[
+            Expansion(token=21, cost=-1.0),
+            Expansion(token=22, cost=-1.0),
+        ])]
     )
     assert expected == expanded
 

run.py

@@ -35,7 +35,10 @@ for i, x in enumerate(contexts):
 batch = Batch(items=series)
 
 #%%
-expanded = expand(batch, expander)
+stopping_criterion = create_stopping_criterion_llm(tokenizer)
+
+#%%
+expanded = expand(batch, expander, stopping_criterion)
 
 # %%
 def print_expansions(expansions: ExpansionResultBatch):