Working generic expand

completions.py

@@ -1,5 +1,6 @@
 #%%
 from dataclasses import dataclass
+import math
 import time
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, BatchEncoding
@@ -34,7 +35,7 @@ def split_into_words(token_probs: list[tuple[int, float]], tokenizer: Tokenizer)
 
     def merge_tokens(a: Tok, b: Tok) -> Tok | None:
         if is_beginning_of_word(a.str) and is_continuation_of_word(b.str):
-            return Tok(b.index, a.ids + b.ids, a.str + b.str, a.logprob * b.logprob)
+            return Tok(a.index, a.ids + b.ids, a.str + b.str, a.logprob + b.logprob)
         return None
 
     converted = [Tok(i, [token_id], tokenizer.decode([token_id]), logprob)
@@ -94,6 +95,20 @@ def generate_outputs(model: PreTrainedModel, inputs: BatchEncoding, num_samples:
         )
     return outputs
 
+def find_next_tokens(model: PreTrainedModel, inputs: BatchEncoding, tokenizer: Tokenizer, min_p: float) -> list[list[tuple[int, str, float]]]:
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    with torch.no_grad():
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    logits: torch.Tensor = outputs.logits[:, -1, :]
+    log_probs: torch.Tensor = torch.log_softmax(logits, dim=-1)
+    # for every batch item, find all tokens with log prob greater than min_p, and return their ids and log probs
+    result = []
+    print(f"{log_probs.shape=}")
+    for probs in log_probs:
+        result.append([(i, tokenizer.convert_ids_to_tokens([i])[0], p) for i, p in enumerate(probs) if p > min_p])
+    return result
+
 def extract_replacements(outputs: GenerateOutput | torch.LongTensor, tokenizer: Tokenizer, num_inputs: int, input_len: int, num_samples: int = 5) -> list[list[str]]:
     all_new_words = []
     for i in range(num_inputs):
@@ -117,30 +132,24 @@ def load_model() -> tuple[PreTrainedModel, Tokenizer, torch.device]:
     return model, tokenizer, device
 
 def check_text(input_text: str, model: PreTrainedModel, tokenizer: Tokenizer, device: torch.device) -> list[ApiWord]:
-#%%
     inputs: BatchEncoding = tokenize(input_text, tokenizer, device)
 
-    #%%
     token_probs: list[tuple[int, float]] = calculate_log_probabilities(model, tokenizer, inputs)
 
-    #%%
     words = split_into_words(token_probs, tokenizer)
     log_prob_threshold = -5.0
     low_prob_words = [(i, word) for i, word in enumerate(words) if word.logprob < log_prob_threshold]
 
-    #%%
     contexts = [word.context for _, word in low_prob_words]
     inputs = prepare_inputs(contexts, tokenizer, device)
     input_ids = inputs["input_ids"]
 
-    #%%
     num_samples = 10
     start_time = time.time()
     outputs = generate_outputs(model, inputs, num_samples)
     end_time = time.time()
     print(f"Total time taken for replacements: {end_time - start_time:.4f} seconds")
 
-    #%%
     replacements = extract_replacements(outputs, tokenizer, input_ids.shape[0], input_ids.shape[1], num_samples)
 
     low_prob_words_with_replacements = { i: (w, r) for (i, w), r in zip(low_prob_words, replacements) }
@@ -152,3 +161,31 @@ def check_text(input_text: str, model: PreTrainedModel, tokenizer: Tokenizer, de
         else:
             result.append(ApiWord(text=word.text, logprob=word.logprob, replacements=[]))
     return result
+
+# %%
+model, tokenizer, device = load_model()
+
+#%%
+input_text = "The quick brown fox jumpz over"
+inputs: BatchEncoding = tokenize(input_text, tokenizer, device)
+
+#%%
+token_probs: list[tuple[int, float]] = calculate_log_probabilities(model, tokenizer, inputs)
+
+#%%
+words = split_into_words(token_probs, tokenizer)
+log_prob_threshold = -5.0
+low_prob_words = [(i, word) for i, word in enumerate(words) if word.logprob < log_prob_threshold]
+
+#%%
+contexts = [word.context for _, word in low_prob_words]
+inputs = prepare_inputs(contexts, tokenizer, device)
+input_ids = inputs["input_ids"]
+
+#%%
+next_tokens = find_next_tokens(model, inputs, tokenizer, min_p=-5)
+
+#%%
+next_tokens
+
+# %%

expand.py

@@ -0,0 +1,84 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Protocol
+
+# import torch
+# from transformers import PreTrainedModel
+# from completions import find_next_tokens, Tokenizer
+
+@dataclass
+class Series:
+    id: int
+    tokens: list[int]
+    budget: float
+
+@dataclass
+class Batch:
+    items: list[Series]
+
+@dataclass
+class ExpansionOne:
+    token: int
+    cost: float
+
+@dataclass
+class ExpansionOneResult:
+    series: Series
+    expansions: list[ExpansionOne]
+
+@dataclass
+class ExpansionOneResultBatch:
+    items: list[ExpansionOneResult]
+
+# 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: ...
+
+@dataclass
+class ExpansionResult:
+    series: Series
+    expansions: list[list[int]]
+
+@dataclass
+class ExpansionResultBatch:
+    items: list[ExpansionResult]
+
+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))
+    return results
+
+def compute_expansions(original_series: list[Series], expanded_series: list[Series]) -> ExpansionResultBatch:
+    # 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
+    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)
+    for s in expanded_series:
+        expanded_series_by_id[s.id].append(s.tokens)
+    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)
+        results.append(expansion_result)
+    return ExpansionResultBatch(items=results)
+
+# A compound operation that we can implement generically, relying on an ExpanderOneBatch
+def expand(batch: Batch, expander: ExpanderOneBatch) -> ExpansionResultBatch:
+    completed_series: list[Series] = []
+    current_batch = batch
+    while len(current_batch.items) > 0:
+        current_batch_items = []
+        expanded = expander.expand(current_batch)
+        for item in expanded.items:
+            if len(item.expansions) == 0:
+                completed_series.append(item.series)
+            else:
+                current_batch_items.extend(compute_new_series(item))
+        current_batch = Batch(items=current_batch_items)
+    return compute_expansions(batch.items, completed_series)

expand_test.py

@@ -0,0 +1,161 @@
+from dataclasses import dataclass
+from expand import Series, ExpanderOneBatch, ExpansionOne, Batch, ExpansionOneResult, ExpansionOneResultBatch, ExpansionResult, ExpansionResultBatch, 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[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]
+
+class HardcodedExpanderOneBatch(ExpanderOneBatch):
+    def expand(self, batch: Batch) -> ExpansionOneResultBatch:
+        result = []
+        for s in batch.items:
+            expansions = expand_series(s)
+            result.append(ExpansionOneResult(series=s, expansions=expansions))
+        return ExpansionOneResultBatch(items=result)
+
+expander = HardcodedExpanderOneBatch()
+
+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=[])]
+    )
+    assert expected == expanded
+
+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=[
+            ExpansionOne(token=21, cost=1.0),
+            ExpansionOne(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 = ExpansionOneResultBatch(
+        items=[ExpansionOneResult(series=s, expansions=[
+            ExpansionOne(token=21, cost=1.0),
+            ExpansionOne(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 = ExpansionOneResultBatch(
+        items=[ExpansionOneResult(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=[
+            ExpansionOne(token=33, cost=1.0),
+            ExpansionOne(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 = ExpansionOneResultBatch(
+        items=[
+            ExpansionOneResult(series=s1, expansions=[
+                ExpansionOne(token=41, cost=1.0),
+                ExpansionOne(token=42, cost=1.0),
+            ]),
+            ExpansionOneResult(series=s2, expansions=[
+                ExpansionOne(token=33, cost=1.0),
+                ExpansionOne(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 == ExpansionResultBatch(items=[
+        ExpansionResult(
+            series=Series(id=0, tokens=[1, 21], budget=1.0),
+            expansions=[
+                [31],
+                [32],
+            ]
+        ),
+        ExpansionResult(
+            series=Series(id=1, tokens=[1, 22], budget=1.0),
+            expansions=[
+                [33],
+                [34],
+            ]
+        ),
+    ])
+
+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 == ExpansionResultBatch(items=[
+        ExpansionResult(
+            series=Series(id=0, tokens=[1, 21], budget=2.0),
+            expansions=[
+                [31, 41],
+                [31, 42],
+                [32, 41],
+            ]
+        ),
+        ExpansionResult(
+            series=Series(id=1, tokens=[1, 22], budget=1.0),
+            expansions=[
+                [33],
+                [34],
+            ]
+        ),
+    ])
+
+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 == ExpansionResultBatch(items=[
+        ExpansionResult(
+            series=Series(id=0, tokens=[1, 21], budget=3.0),
+            expansions=[
+                [31, 41],
+                [31, 42],
+                [32, 41, 51],
+            ]
+        ),
+        ExpansionResult(
+            series=Series(id=1, tokens=[1, 22], budget=0.0),
+            expansions=[],
+        ),
+    ])