println profiling

expand.py

@@ -86,7 +86,8 @@ def expand(batch: Batch, expander: BatchExpander, completion_criterion: Callable
     completed_series: list[Series] = []
     current_batch = batch
     while len(current_batch.items) > 0:
-        print(f"Expanding {len(current_batch.items)} series: {current_batch.items}")
+        # print(f"Expanding {len(current_batch.items)} series: {current_batch.items}")
+        print(f"Expanding {len(current_batch.items)} series")
         current_batch_items = []
         start_time = time.time()
         expanded = expander.expand(current_batch)

expand_llm.py

@@ -2,19 +2,28 @@ import torch
 from expand import *
 from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, BatchEncoding
 from dataclasses import dataclass
+import time
 
 type Tokenizer = PreTrainedTokenizer | PreTrainedTokenizerFast
 
 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"]
+    print("Running inference")
+    start_time = time.time()
     with torch.no_grad():
         outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    print(f"Inference done, took {time.time() - start_time} seconds")
+    start_time = time.time()
     logits: torch.Tensor = outputs.logits[:, -1, :]
     log_probs: torch.Tensor = torch.log_softmax(logits, dim=-1)
+    print(f"Log probs done, took {time.time() - start_time} seconds")
+    start_time = time.time()
     result = []
+    print(f"Resulting tensor: {log_probs.shape}")
     for probs in log_probs:
         result.append([(i, p.item()) for i, p in enumerate(probs)])
+    print(f"Result done, took {time.time() - start_time} seconds")
     return result
 
 def prepare_inputs(contexts: list[list[int]], tokenizer: Tokenizer, device: torch.device) -> BatchEncoding:
@@ -29,10 +38,15 @@ class LLMBatchExpander(BatchExpander):
     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)
+        start_time = time.time()
         results = []
+        print(f"Batch size: {len(batch.items)}, next tokens size: {len(next_tokens)}")
         for s, next_tokens in zip(batch.items, next_tokens):
+            print(f"Series {s.id}, {len(next_tokens)=}")
             expansions = [Expansion(token=token, cost=cost) for token, cost in next_tokens]
             results.append(TokenCandidates(series=s, expansions=expansions))
+        print()
+        print(f"Token candidates done, took {time.time() - start_time} seconds")
         return BatchCandidates(items=results)
 
 def create_stopping_criterion_llm(tokenizer: Tokenizer) -> Callable[[Series, Expansion], bool]:
@@ -42,7 +56,7 @@ def create_stopping_criterion_llm(tokenizer: Tokenizer) -> Callable[[Series, Exp
             return d
         token_str = tokenizer.decode([expansion.token])
         starts_with_space = token_str.startswith(" ")
-        print(f"-----{token_str}-----, {starts_with_space=}")
+        # 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