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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.1.6
+Version: 0.1.7
 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.6 → divergent_beamsearch-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "divergent-beamsearch"
-version = "0.1.6"
+version = "0.1.7"
 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.6 → divergent_beamsearch-0.1.7}/src/divergent_beamsearch/algorithm.py

@@ -76,6 +76,12 @@ class AcceptEverythingParser(Parser):
     def copy(self):
         return self
 
+def index_reduce_lists(x : torch.Tensor, indices : list[list[int]], reduce_func=torch.sum) -> torch.Tensor:
+    values = []
+    for i, index in enumerate(indices):
+        values.append(reduce_func(x[i, index], dim=-1))
+    return torch.tensor(values, dtype=x.dtype, device=x.device, requires_grad=x.requires_grad)
+
 @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]:
     assert input_ids.shape[0] == 1, "Batch size must be 1"
@@ -120,7 +126,8 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
 
         scores_finished_current = scores_unfinished[can_end]
         solutions_finished_current = solutions_unfinished[can_end]
-        scores_finished_current = scores_finished_current + log1mexp(logprobs[can_end, select_mask(parsers_tokens, can_end)].logsumexp(dim=-1)).squeeze(-1)
+        logprob_other_ans = index_reduce_lists(logprobs[can_end], select_mask(parsers_tokens, can_end), reduce_func=torch.logsumexp).squeeze(-1)
+        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)

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

@@ -37,10 +37,10 @@ def fakemodel_and_tokenizer():
 
 @pytest.mark.parametrize("device", ['cpu', 'cuda'])
 @pytest.mark.parametrize("end_symb", TEST_END_SYMBS)
-def test_divergent_beamsearch(fakemodel_and_tokenizer, device, end_symb):
+def test_divergent_beamsearch(model_and_tokenizer, device, end_symb):
     if device == 'cuda' and not torch.cuda.is_available():
         pytest.skip("CUDA is not available on this machine.")
-    model, tokenizer = fakemodel_and_tokenizer
+    model, tokenizer = model_and_tokenizer
     model.to(device)
     prompt = "The capital of France is"
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
@@ -48,7 +48,7 @@ def test_divergent_beamsearch(fakemodel_and_tokenizer, device, end_symb):
     max_length = 10
     pad_token_id = tokenizer.eos_token_id
 
-    possible_answers = [' Paris', ' Paris Hilton']
+    possible_answers = [' Paris', ' Madrid', ' Paris Hilton']
     tokenized_answers = tokenizer(possible_answers).input_ids
 
     if end_symb == 'tokenizer':
@@ -56,9 +56,12 @@ def test_divergent_beamsearch(fakemodel_and_tokenizer, device, end_symb):
     
     multi_choices_parser = MultiChoicesParser([tokenized_answers], end_symb=end_symb)
 
-    logprob_paris = model(input_ids).logits.cpu().log_softmax(dim=-1)[0, -1, tokenized_answers[0][0]]
-    logprob_hilton = model(torch.cat([input_ids, torch.tensor(tokenized_answers[1][0], device=device).view(1,1)], dim=-1)).logits.cpu().log_softmax(dim=-1)[0, -1, tokenized_answers[1][1]]
-    logprob_paris_hilton = logprob_paris + logprob_hilton
+    with torch.no_grad():
+        logprob_paris = model(input_ids).logits.cpu().log_softmax(dim=-1)[0, -1, tokenized_answers[0][0]]
+        logprob_hilton = model(torch.cat([input_ids, torch.tensor(tokenized_answers[2][0], device=device).view(1,1)], dim=-1)).logits.cpu().log_softmax(dim=-1)[0, -1, tokenized_answers[2][1]]
+        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)
 
     scores, solutions = divergent_beamsearch(
         input_ids=input_ids,
@@ -72,8 +75,10 @@ def test_divergent_beamsearch(fakemodel_and_tokenizer, device, 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)
     assert (solutions == true_solutions).all(), "Beam search did not return the expected solutions"
-    assert scores[0] == logprob_paris + log1mexp(logprob_hilton), "Beam search did not return the expected score"
-    assert scores[1] == logprob_paris_hilton, "Beam search did not return the expected score"
+    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"
+
 
 @pytest.mark.parametrize("device", ['cpu', 'cuda'])
 @pytest.mark.parametrize("end_symb", TEST_END_SYMBS)