divergent-beamsearch 0.1.8__tar.gz → 0.2.1__tar.gz

{divergent_beamsearch-0.1.8 → divergent_beamsearch-0.2.1}/PKG-INFO

@@ -1,12 +1,11 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.1.8
+Version: 0.2.1
 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.9.72
 Requires-Dist: torch>=2.0.0
-Requires-Dist: transformers>=4.47.1
 Description-Content-Type: text/markdown
 
 # Divergent Beam Search

{divergent_beamsearch-0.1.8 → divergent_beamsearch-0.2.1}/pyproject.toml

@@ -1,18 +1,18 @@
 [project]
 name = "divergent-beamsearch"
-version = "0.1.8"
+version = "0.2.1"
 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.9.72",
     "torch>=2.0.0",
-    "transformers>=4.47.1",
 ]
 
 [dependency-groups]
 dev = [
     "pytest>=8.3.4",
+    "transformers>=4.47.1"
 ]
 
 [build-system]

{divergent_beamsearch-0.1.8 → divergent_beamsearch-0.2.1}/src/divergent_beamsearch/algorithm.py

@@ -1,6 +1,9 @@
 import math
 import torch
-from transformers import GPT2LMHeadModel
+try:
+    from transformers import GPT2LMHeadModel
+except ImportError:
+    pass
 from multi_choices_parser import DEFAULT_END_SYMB
 
 
@@ -35,12 +38,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:
@@ -91,7 +99,9 @@ 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, 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()
@@ -114,7 +124,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)
@@ -173,21 +183,22 @@ 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) -> 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.8 → divergent_beamsearch-0.2.1}/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
 
@@ -77,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'])
@@ -95,7 +99,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)
@@ -200,3 +203,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()

{divergent_beamsearch-0.1.8 → divergent_beamsearch-0.2.1}/uv.lock

@@ -73,7 +73,7 @@ wheels = [
 
 [[package]]
 name = "divergent-beamsearch"
-version = "0.1.5"
+version = "0.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "multi-choices-parser" },
@@ -88,7 +88,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "multi-choices-parser", specifier = ">=0.9.61" },
+    { name = "multi-choices-parser", specifier = ">=0.9.72" },
     { name = "torch", specifier = ">=2.0.0" },
     { name = "transformers", specifier = ">=4.47.1" },
 ]
@@ -221,11 +221,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.9.72"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ab/14/7a99908c455ed355563c1a59c3953fd2e1e0b8bd3699f616adf44f31c019/multi_choices_parser-0.9.72-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7554b9928b663952d50dad2be070b33eac12a7cf0a5d0237ca273f075e598d09", size = 99165 },
+    { url = "https://files.pythonhosted.org/packages/aa/3c/fb13affb1061050fb0f2988d1fdd0f37943e17abf1644ac681d6cda45615/multi_choices_parser-0.9.72-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4700eab86215bd4f0da9fcea0650e5336bc22d62a77625b2d3d1b1a83081b0d4", size = 139545 },
+    { url = "https://files.pythonhosted.org/packages/47/c0/5b47daed1dd6cff64c602cdcefda285eacfdf71f43d0452ed2f68e17ae9e/multi_choices_parser-0.9.72-cp311-cp311-win_amd64.whl", hash = "sha256:65725c593363b8c207748478ca966e5fc0288118b95c2e6b7cc338003417a185", size = 105544 },
+    { url = "https://files.pythonhosted.org/packages/51/98/10331d2da4c0c036720f1cd41a60f33cf35a4ac2aad963dd58e486d97ccb/multi_choices_parser-0.9.72-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:277ed1e6e4c7044313281caa5c20cb09eab518f27271001afea748793acc26a9", size = 99360 },
+    { url = "https://files.pythonhosted.org/packages/98/6a/5c90c3b19013aa02b40ed6ef193213bfb4ad92e4500c8e1009e712c0d6db/multi_choices_parser-0.9.72-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:441ef6dd784c9d7fbf0effe66f2e910ad308604749924e0dccec79fd24cfdf2e", size = 138117 },
+    { url = "https://files.pythonhosted.org/packages/07/13/4c601c9336b7a83e762937c2d75823964a9a9773903cba2696ec59107dbf/multi_choices_parser-0.9.72-cp312-cp312-win_amd64.whl", hash = "sha256:28ac8cea47639b434fc88e143f5d38a0bd5ab4ce9a040a036e532896185d672b", size = 105658 },
+    { url = "https://files.pythonhosted.org/packages/db/9d/8ea1f8a87282da07b0d5044c682566a68eced933bb675c8936400bb72a54/multi_choices_parser-0.9.72-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3dfd87c968a6e5618a0cacc3d29fa244cf427208829eeda82802fd60250ea1b3", size = 99398 },
+    { url = "https://files.pythonhosted.org/packages/8b/8d/f0a244c59e13e4591e5be9f0793a22e3cde6b631801f9473e96fe44c76bb/multi_choices_parser-0.9.72-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c99dae2ba228c9362648ed66bd9790df07a71176f22938595487605b489a8dc8", size = 138403 },
+    { url = "https://files.pythonhosted.org/packages/f3/f2/73e929b894fe379be5e1b8d373c9acf66c5b34da3f886edbe22ff8725593/multi_choices_parser-0.9.72-cp313-cp313-win_amd64.whl", hash = "sha256:9576300f71ba688f799832e8b86b3cb24ea74cde29aa4e70ac63ec7545e32790", size = 105658 },
 ]
 
 [[package]]