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

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

@@ -1,12 +1,11 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.2.0
+Version: 0.2.2
 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
-Requires-Dist: multi-choices-parser>=0.9.61
+Requires-Dist: multi-choices-parser>=0.10.0
 Requires-Dist: torch>=2.0.0
-Requires-Dist: transformers>=4.47.1
 Description-Content-Type: text/markdown
 
 # Divergent Beam Search

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

@@ -1,18 +1,18 @@
 [project]
 name = "divergent-beamsearch"
-version = "0.2.0"
+version = "0.2.2"
 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"
 dependencies = [
-    "multi-choices-parser>=0.9.61",
+    "multi-choices-parser>=0.10.0",
     "torch>=2.0.0",
-    "transformers>=4.47.1",
 ]
 
 [dependency-groups]
 dev = [
     "pytest>=8.3.4",
+    "transformers>=4.47.1"
 ]
 
 [build-system]

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

@@ -1,6 +1,10 @@
 import math
+import multi_choices_parser
 import torch
-from transformers import GPT2LMHeadModel
+try:
+    from transformers import GPT2LMHeadModel
+except ImportError:
+    pass
 from multi_choices_parser import DEFAULT_END_SYMB
 
 
@@ -35,7 +39,7 @@ 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, 
+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 = []
@@ -96,7 +100,7 @@ def pad_to_same_size(tensors : list[torch.Tensor], padding_value : int) -> torch
     return torch.cat(padded_tensors, dim=0)
 
 @torch.no_grad()
-def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam_size : int, 
+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"
@@ -160,9 +164,11 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
         input_ids_unfinished = torch.cat([input_ids_unfinished[best_tokens_row], best_tokens.unsqueeze(-1)], dim=-1)
         scores_unfinished = scores_unfinished[best_tokens_row] + best_tokens_logprobs
         solutions_unfinished = torch.cat([solutions_unfinished[best_tokens_row], best_tokens.unsqueeze(-1)], dim=-1)
+        best_tokens_row = best_tokens_row.tolist()
         parsers_unfinished = [parsers_unfinished[row].copy() for row in best_tokens_row]
-        for parser, token in zip(parsers_unfinished, best_tokens.tolist()):
-            parser.step(token)
+        for parser, token, row in zip(parsers_unfinished, best_tokens.tolist(), best_tokens_row):
+            if not parser.finished:
+                parser.step(token)
 
     # Special case of vanilla beam search where all answers are valid
     # Warning : In this case model will not stop on end_of_sentence token
@@ -180,10 +186,11 @@ def set_slice_row(x : torch.Tensor, slices : torch.IntTensor, value) -> torch.Te
         x[i].index_fill_(0, indices[i], 0)
 
 @torch.no_grad()
-def divergent_logprob(input_ids : torch.Tensor, attention_mask : torch.Tensor | None, model : GPT2LMHeadModel, 
+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, optimize_gpu_mem=True) -> torch.FloatTensor:
     if start is None:
+        # Start at 1 because first token logprobs cannot be computed
         start = 1
     if isinstance(start, int):
         start = torch.tensor([start]*input_ids.shape[0])
@@ -222,8 +229,9 @@ def divergent_logprob(input_ids : torch.Tensor, attention_mask : torch.Tensor |
         for input_id, att in zip(input_ids[i, start:].tolist(), attention_mask[i, start:].tolist()):
             if not att:
                 break
+            assert not parser.finished
             parser.step(input_id)
-        next_tokens = list(parser.next())
+        next_tokens = parser.next()
         try:
             next_tokens.remove(end_symb)
         except ValueError:

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

@@ -79,11 +79,13 @@ def test_divergent_beamsearch(model_and_tokenizer, device, end_symb):
         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)
-    assert (solutions == true_solutions).all(), "Beam search did not return the expected solutions"
+    
     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"
+    assert (solutions == true_solutions).all(), "Beam search did not return the expected solutions"
+
 
 
 @pytest.mark.parametrize("device", ['cpu', 'cuda'])
@@ -215,9 +217,8 @@ def test_element_wise_equivalence_divergent_logprob(fakemodel_and_tokenizer, dev
         'Google is owned by Alphabet'
     ]
     
-    multi_choices_parser = MultiChoicesParser([texts])
-
     inputs = tokenizer(texts, return_tensors='pt', padding=True).to(device)
+    multi_choices_parser = MultiChoicesParser([[x[1:] for x in tokenizer(texts).input_ids]])
 
     logprobs_global = divergent_logprob(inputs.input_ids, inputs.attention_mask, model, multi_choices_parser)
 

{divergent_beamsearch-0.2.0 → divergent_beamsearch-0.2.2}/uv.lock

@@ -73,28 +73,30 @@ wheels = [
 
 [[package]]
 name = "divergent-beamsearch"
-version = "0.1.5"
+version = "0.2.2"
 source = { editable = "." }
 dependencies = [
     { name = "multi-choices-parser" },
     { name = "torch" },
-    { name = "transformers" },
 ]
 
 [package.dev-dependencies]
 dev = [
     { name = "pytest" },
+    { name = "transformers" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "multi-choices-parser", specifier = ">=0.9.61" },
+    { name = "multi-choices-parser", specifier = ">=0.10.0" },
     { name = "torch", specifier = ">=2.0.0" },
-    { name = "transformers", specifier = ">=4.47.1" },
 ]
 
 [package.metadata.requires-dev]
-dev = [{ name = "pytest", specifier = ">=8.3.4" }]
+dev = [
+    { name = "pytest", specifier = ">=8.3.4" },
+    { name = "transformers", specifier = ">=4.47.1" },
+]
 
 [[package]]
 name = "filelock"
@@ -221,11 +223,18 @@ wheels = [
 
 [[package]]
 name = "multi-choices-parser"
-version = "0.9.61"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/e2/17/90a6125a2145c03e39c3e7f78f65121eb14dedb9de8b40aee3c8a24a709b/multi_choices_parser-0.9.61.tar.gz", hash = "sha256:be932cac4aeabe9ee057c6d4592ea4325a0a92e52758d77a9e08bafa2cd23294", size = 7889 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/3c/4f/c5a514a510779202ff37505220edfba9154ceff31958ed71fa1878781af9/multi_choices_parser-0.9.61-py3-none-any.whl", hash = "sha256:36bc367bceb66bbfb1bea26d9a38aa9cd10273b54cef331dd7c69da582fb9c2a", size = 6870 },
+version = "0.10.0"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/25/59/233da6ab703cf3243dffd2180d082a91a45caf720723309090cee3353da7/multi_choices_parser-0.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3019671e9ed6daa0fb8c746cd9f52557da280d57a3ba938f72a34db336671980", size = 99282 },
+    { url = "https://files.pythonhosted.org/packages/30/b0/82b5ea3ebb500df180cf15e2d7d43bbcef1d58b122206f0b4616bf1dabf5/multi_choices_parser-0.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bebb99bf5096e8f40ee584e5860efa10e9484a50c7747360f313c761e16ed5c5", size = 139668 },
+    { url = "https://files.pythonhosted.org/packages/2d/f4/b7e12764e7366b01d7fa5fdd177480967492219b1d7ffd5c6a35f8117247/multi_choices_parser-0.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:bee4d14b626fa9f8290670047bc0cc358c0a3dddd0dc104e9e844ed4b1b43203", size = 105669 },
+    { url = "https://files.pythonhosted.org/packages/d6/08/f6eacab1476d99b64443433e5a683afaf79f8ae6798edf12a7535a7a02af/multi_choices_parser-0.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:818d292bfa2d35e0cbee4608cb2f4f9e223c7a8b8d3f94a83dc83d05f4dd71ef", size = 99485 },
+    { url = "https://files.pythonhosted.org/packages/02/6a/6ecfcf3b14972807cf3eb34d960691116bbdd655ba5466905543b0fb0a53/multi_choices_parser-0.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4e5ed7468834cf9502d3aa3fe71eaf9b2a28e3e40eb60744ded7a7605eed3612", size = 138236 },
+    { url = "https://files.pythonhosted.org/packages/62/04/7ab6935f99d275fba202cd21b2cd0fb2f775237c6b57ad247cbb95e4db53/multi_choices_parser-0.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:fbc8c4b6f8bbf9e2ead1c228b6a5be9fac4d91854797b430fcb05d91ba96f8dd", size = 105784 },
+    { url = "https://files.pythonhosted.org/packages/27/88/8fb06ff9341e4a09b714939d515e583678e1620f2d3a4536e4776a4ad92a/multi_choices_parser-0.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:541f5ff7da3cb7fb1b4ef2d114ee02637322c058c84345283c9c3194c7207e31", size = 99523 },
+    { url = "https://files.pythonhosted.org/packages/e6/da/7f853bb1e676d74c85d25a1023674bcb0407d9a222ce9f65d56de4025dd5/multi_choices_parser-0.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ead4849057c50608a48eb498a95a622a2b0151e20c871ed9ec27ed27eb20108d", size = 138524 },
+    { url = "https://files.pythonhosted.org/packages/ca/2f/7b8baffc032b503fc1075fa0be19c8ab3b56265b8c3a763bfac6c27b835f/multi_choices_parser-0.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:ad17159761164672895efe07c75e1d872c8c40e74e89aa9cb6ff74fd9a81362d", size = 105781 },
 ]
 
 [[package]]