divergent-beamsearch 0.1.7__tar.gz → 0.2.0__tar.gz

{divergent_beamsearch-0.1.7 → divergent_beamsearch-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.1.7
+Version: 0.2.0
 Summary: A variant of the beam search algorithm that focuses on finding answers that maximize the probability of generating an answer before diverging into another subject.
 License-File: LICENCE
 Requires-Python: >=3.11

{divergent_beamsearch-0.1.7 → divergent_beamsearch-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "divergent-beamsearch"
-version = "0.1.7"
+version = "0.2.0"
 description = "A variant of the beam search algorithm that focuses on finding answers that maximize the probability of generating an answer before diverging into another subject."
 readme = "README.md"
 requires-python = ">=3.11"

{divergent_beamsearch-0.1.7 → divergent_beamsearch-0.2.0}/src/divergent_beamsearch/algorithm.py

@@ -35,12 +35,17 @@ def apply_mask_tokens(pred : torch.Tensor, parsers_tokens):
     return pred[~pred.isinf().all(dim=-1)]
 
 
-def batched_inference_logits(model : GPT2LMHeadModel, input_ids : torch.Tensor, attention_mask : torch.Tensor | None = None, batch_size : int = 32) -> torch.Tensor:
+def batched_inference_logits(model : GPT2LMHeadModel, input_ids : torch.Tensor, 
+                             attention_mask : torch.Tensor | None = None, batch_size : int = 32,
+                             to_cpu=False) -> torch.Tensor:
     logits = []
     if attention_mask is None:
         attention_mask = torch.ones_like(input_ids)
     for i in range(0, input_ids.shape[0], batch_size):
-        logits.append(model(input_ids[i:i+batch_size], attention_mask=attention_mask[i:i+batch_size]).logits)
+        l = model(input_ids[i:i+batch_size], attention_mask=attention_mask[i:i+batch_size]).logits
+        if to_cpu:
+            l = l.cpu()
+        logits.append(l)
     return torch.cat(logits, dim=0)
 
 def select_mask(source : list, mask : list[bool]) -> list:
@@ -82,8 +87,18 @@ def index_reduce_lists(x : torch.Tensor, indices : list[list[int]], reduce_func=
         values.append(reduce_func(x[i, index], dim=-1))
     return torch.tensor(values, dtype=x.dtype, device=x.device, requires_grad=x.requires_grad)
 
+def pad_to_same_size(tensors : list[torch.Tensor], padding_value : int) -> torch.Tensor:
+    max_size = max(x.shape[-1] for x in tensors)
+    padded_tensors = []
+    for tensor in tensors:
+        pad = torch.full((tensor.shape[0], max_size - tensor.shape[1]), padding_value, dtype=torch.long)
+        padded_tensors.append(torch.cat([tensor, pad], dim=-1))
+    return torch.cat(padded_tensors, dim=0)
+
 @torch.no_grad()
-def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam_size : int, max_length : int, parser : Parser, pad_token_id : int, batch_size=32, num_solutions = None, end_symb=DEFAULT_END_SYMB) -> tuple[torch.Tensor, torch.Tensor]:
+def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam_size : int, 
+                         max_length : int, parser : Parser, pad_token_id : int, batch_size=32, 
+                         num_solutions = None, end_symb=DEFAULT_END_SYMB, optimize_gpu_mem=True) -> tuple[torch.Tensor, torch.Tensor]:
     assert input_ids.shape[0] == 1, "Batch size must be 1"
     device = input_ids.device
     input_ids = input_ids.cpu()
@@ -106,7 +121,7 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
     for _ in range(max_length):
         if len(input_ids_unfinished) == 0:
             break
-        pred = batched_inference_logits(model, input_ids_unfinished.to(device), batch_size=batch_size)[:, -1].cpu()
+        pred = batched_inference_logits(model, input_ids_unfinished.to(device), batch_size=batch_size, to_cpu=optimize_gpu_mem)[:, -1].cpu()
         parsers_tokens, can_end = get_parsers_tokens(parsers_unfinished, end_symb)
         logprobs = torch.log_softmax(pred, dim=-1)
         logprobs_filtered = apply_mask_tokens(logprobs, parsers_tokens)
@@ -130,8 +145,11 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
         scores_finished_current = scores_finished_current + log1mexp(logprob_other_ans)
         scores_finished = torch.cat([scores_finished, scores_finished_current])
         if len(solutions_finished_current):
-            pad = torch.full((len(scores_finished_current), solutions_finished_current.shape[1] - solutions_finished.shape[1]), pad_token_id, dtype=torch.long)
-            solutions_finished = torch.cat([solutions_finished.view(-1, solutions_finished_current.shape[1]+pad.shape[1]), torch.cat([solutions_finished_current, pad], dim=1)], dim=0)
+            if len(solutions_finished):
+                solutions_finished = pad_to_same_size([solutions_finished, solutions_finished_current], 
+                                                                    padding_value=pad_token_id)
+            else:
+                solutions_finished = solutions_finished_current
         if solutions_finished.numel():
             # Keep num_solutions best solutions in finished
             order = scores_finished.argsort(descending=True)
@@ -164,19 +182,20 @@ def set_slice_row(x : torch.Tensor, slices : torch.IntTensor, value) -> torch.Te
 @torch.no_grad()
 def divergent_logprob(input_ids : torch.Tensor, attention_mask : torch.Tensor | None, model : GPT2LMHeadModel, 
                       parsers : Parser | list[Parser] | None, batch_size=32, 
-                      start : int | torch.IntTensor = None, end_symb=DEFAULT_END_SYMB) -> torch.FloatTensor:
+                      start : int | torch.IntTensor = None, end_symb=DEFAULT_END_SYMB, optimize_gpu_mem=True) -> torch.FloatTensor:
     if start is None:
-        start = 0
+        start = 1
     if isinstance(start, int):
         start = torch.tensor([start]*input_ids.shape[0])
     assert start.shape[0] == input_ids.shape[0]
+    assert (start > 0).all()
     # -1 because next token offset
     start = start - 1
 
     if attention_mask is None:
         attention_mask = torch.ones_like(input_ids)
 
-    logits = batched_inference_logits(model, input_ids, attention_mask, batch_size).cpu()
+    logits = batched_inference_logits(model, input_ids, attention_mask, batch_size, to_cpu=optimize_gpu_mem).cpu()
     input_ids = input_ids.cpu()
     attention_mask = attention_mask.cpu()
 

{divergent_beamsearch-0.1.7 → divergent_beamsearch-0.2.0}/tests/test_beamsearch.py

@@ -13,6 +13,7 @@ TEST_END_SYMBS = [DEFAULT_END_SYMB, 'tokenizer']
 def model_and_tokenizer():
     model = GPT2LMHeadModel.from_pretrained("gpt2")
     tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
     return model, tokenizer
 
 @pytest.fixture
@@ -32,6 +33,7 @@ def fakemodel_and_tokenizer():
     # Instantiate a model with the custom configuration
     model = GPT2LMHeadModel(config)
     model.eval()
+    tokenizer.pad_token = tokenizer.eos_token
 
     return model, tokenizer
 
@@ -44,11 +46,11 @@ def test_divergent_beamsearch(model_and_tokenizer, device, end_symb):
     model.to(device)
     prompt = "The capital of France is"
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
-    beam_size = 5
+    beam_size = 10
     max_length = 10
     pad_token_id = tokenizer.eos_token_id
 
-    possible_answers = [' Paris', ' Madrid', ' Paris Hilton']
+    possible_answers = [' Paris', ' Madrid', ' Paris Hilton', ' Bri bra brouuu Mario Brooos']
     tokenized_answers = tokenizer(possible_answers).input_ids
 
     if end_symb == 'tokenizer':
@@ -62,6 +64,9 @@ def test_divergent_beamsearch(model_and_tokenizer, device, end_symb):
         logprob_paris_hilton = logprob_paris + logprob_hilton
         logprob_madrid = model(input_ids).logits.cpu().log_softmax(dim=-1)[0, -1, tokenized_answers[1][0]]
         logprob_paris_diverge = logprob_paris + log1mexp(logprob_hilton)
+        input_garbage = torch.tensor(input_ids.tolist()[0] + tokenized_answers[-1]).unsqueeze(0).to(device)
+        logsoftmax_garbage = model(input_garbage).logits.log_softmax(-1)
+        logprob_garbage = torch.gather(logsoftmax_garbage[:, 4:-1, :], 2, input_garbage[:, 5:, None]).squeeze(-1).sum(-1)
 
     scores, solutions = divergent_beamsearch(
         input_ids=input_ids,
@@ -70,7 +75,7 @@ def test_divergent_beamsearch(model_and_tokenizer, device, end_symb):
         max_length=max_length,
         parser=multi_choices_parser,
         pad_token_id=pad_token_id,
-        num_solutions=10,
+        num_solutions=beam_size,
         end_symb=end_symb
     )
     true_solutions = torch.nn.utils.rnn.pad_sequence([torch.tensor(ans) for ans in tokenized_answers], batch_first=True, padding_value=pad_token_id)
@@ -78,6 +83,7 @@ def test_divergent_beamsearch(model_and_tokenizer, device, end_symb):
     assert torch.isclose(scores[0], logprob_paris_diverge), "Beam search did not return the expected score"
     assert torch.isclose(scores[1], logprob_madrid), "Beam search did not return the expected score"
     assert torch.isclose(scores[2], logprob_paris_hilton), "Beam search did not return the expected score"
+    assert torch.isclose(scores[3], logprob_garbage), "Beam search did not return the expected score"
 
 
 @pytest.mark.parametrize("device", ['cpu', 'cuda'])
@@ -91,7 +97,6 @@ def test_divergent_logprob(fakemodel_and_tokenizer, device, end_symb):
         "The capital of France is Paris",
         "The top model Paris Hilton"
     ]
-    tokenizer.pad_token = tokenizer.eos_token
     inp = tokenizer(prompts, return_tensors="pt", padding=True)
     input_ids = inp.input_ids.to(device)
     attention_mask = inp.attention_mask.to(device)
@@ -196,3 +201,32 @@ def test_vanilla_beamsearch(model_and_tokenizer, device):
     assert np.isclose(
         scores.cpu().numpy(), np.array([-8.1361, -8.7745, -9.1053]), atol=0.0001
     ).all()
+
+@pytest.mark.parametrize("device", ['cpu', 'cuda'])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+def test_element_wise_equivalence_divergent_logprob(fakemodel_and_tokenizer, device, dtype):
+    model, tokenizer = fakemodel_and_tokenizer
+    model.to(device)
+    model.to(dtype)
+
+    texts = [
+        'My name is Roger',
+        'The capital of Morocco is Rabat',
+        'Google is owned by Alphabet'
+    ]
+    
+    multi_choices_parser = MultiChoicesParser([texts])
+
+    inputs = tokenizer(texts, return_tensors='pt', padding=True).to(device)
+
+    logprobs_global = divergent_logprob(inputs.input_ids, inputs.attention_mask, model, multi_choices_parser)
+
+    logprobs_individual = []
+
+    for text in texts:
+        inputs = tokenizer(text, return_tensors='pt', padding=True).to(device)
+        input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
+        logprobs_individual.append(divergent_logprob(input_ids, attention_mask, model, multi_choices_parser))
+    logprobs_individual = torch.tensor(logprobs_global)
+
+    assert (logprobs_individual == logprobs_global).all()