Working expand_llm

completions.py

@@ -95,7 +95,7 @@ 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]]]:
+def find_next_tokens_0(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():
@@ -109,6 +109,18 @@ def find_next_tokens(model: PreTrainedModel, inputs: BatchEncoding, tokenizer: T
         result.append([(i, tokenizer.convert_ids_to_tokens([i])[0], p) for i, p in enumerate(probs) if p > min_p])
     return result
 
+def find_next_tokens(model: PreTrainedModel, inputs: BatchEncoding, tokenizer: Tokenizer) -> list[list[tuple[int, 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)
+    result = []
+    for probs in log_probs:
+        result.append([(i, p) for i, p in enumerate(probs)])
+    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):
@@ -161,31 +173,3 @@ 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

@@ -2,10 +2,6 @@ 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
@@ -46,7 +42,7 @@ 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 + [expansion.token], budget=result.series.budget + expansion.cost))
     return results
 
 def compute_expansions(original_series: list[Series], expanded_series: list[Series]) -> ExpansionResultBatch:
@@ -73,6 +69,7 @@ def expand(batch: Batch, expander: ExpanderOneBatch) -> ExpansionResultBatch:
     completed_series: list[Series] = []
     current_batch = batch
     while len(current_batch.items) > 0:
+        print(f"Expanding {len(current_batch.items)} series: {current_batch.items}")
         current_batch_items = []
         expanded = expander.expand(current_batch)
         for item in expanded.items:

expand_llm.py

@@ -0,0 +1,20 @@
+from expand import *
+from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, BatchEncoding
+from dataclasses import dataclass
+from completions import prepare_inputs, find_next_tokens
+
+type Tokenizer = PreTrainedTokenizer | PreTrainedTokenizerFast
+
+@dataclass
+class ExpanderOneBatchLLM:
+    model: PreTrainedModel
+    tokenizer: Tokenizer
+
+    def expand(self, batch: Batch) -> ExpansionOneResultBatch:
+        inputs = prepare_inputs([s.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 = [ExpansionOne(token=token, cost=logprob) for token, logprob in next_tokens if logprob + s.budget >= 0]
+            results.append(ExpansionOneResult(series=s, expansions=expansions))
+        return ExpansionOneResultBatch(items=results)

expand_test.py

@@ -12,8 +12,8 @@ possible_sequences = [
 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]
+    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]
 
 class HardcodedExpanderOneBatch(ExpanderOneBatch):
     def expand(self, batch: Batch) -> ExpansionOneResultBatch:
@@ -38,8 +38,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),
+            ExpansionOne(token=21, cost=-1.0),
+            ExpansionOne(token=22, cost=-1.0),
         ])]
     )
     assert expected == expanded
@@ -49,8 +49,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),
+            ExpansionOne(token=21, cost=-1.0),
+            ExpansionOne(token=22, cost=-1.0),
         ])]
     )
     assert expected == expanded
@@ -68,8 +68,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),
+            ExpansionOne(token=33, cost=-1.0),
+            ExpansionOne(token=34, cost=-1.0),
         ])]
     )
     assert expected == expanded
@@ -81,12 +81,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),
+                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),
+                ExpansionOne(token=33, cost=-1.0),
+                ExpansionOne(token=34, cost=-1.0),
             ])
         ]
     )

run.py

@@ -0,0 +1,49 @@
+#%%
+from completions import *
+from expand_llm import *
+from expand import *
+
+# %%
+model, tokenizer, device = load_model()
+
+#%%
+# input_text = "The quick brown fox jumpz over"
+# input_text = "He asked me to prostate myself before the king"
+input_text = "Здравствуйте, я хочу предвыполнить заказ"
+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]
+
+#%%
+expander = ExpanderOneBatchLLM(model, tokenizer)
+
+#%%
+series = []
+for i, x in enumerate(contexts):
+    series.append(Series(id=i, tokens=x, budget=5.0))
+
+#%%
+batch = Batch(items=series)
+
+#%%
+expanded = expand(batch, expander)
+
+# %%
+def print_expansions(expansions: ExpansionResultBatch):
+    for result in expansions.items:
+        for expansion in result.expansions:
+            # convert tokens to string
+            s = tokenizer.decode(expansion)
+            print(f"{result.series.id}: {expansion} {s}")
+
+print_expansions(expanded)
+# %%