divergent-beamsearch 0.1.0__tar.gz → 0.1.1__tar.gz

{divergent_beamsearch-0.1.0 → divergent_beamsearch-0.1.1}/.gitignore

@@ -1,4 +1,5 @@
 **/__pycache__
 .pytest_cache
 .vscode
-.venv
+.venv
+dist

{divergent_beamsearch-0.1.0 → divergent_beamsearch-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.1.0
+Version: 0.1.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

{divergent_beamsearch-0.1.0 → divergent_beamsearch-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "divergent-beamsearch"
-version = "0.1.0"
+version = "0.1.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"

{divergent_beamsearch-0.1.0 → divergent_beamsearch-0.1.1}/src/divergent_beamsearch/algorithm.py

@@ -47,12 +47,29 @@ def log1mexp(x: torch.Tensor) -> torch.Tensor:
         (-x.exp()).log1p(),
     )
 
+class AcceptEverythingParser:
+    def __init__(self, vocab_size : int):
+        self.vocab_size = vocab_size
+        self.tokens = tuple(range(vocab_size))
+
+    def step(self, token):
+        pass
+
+    def next(self):
+        return self.tokens
+    
+    def copy(self):
+        return self
+
 @torch.no_grad()
 def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam_size : int, max_length : int, multi_choices_parser : MultiChoicesParser, pad_token_id : int, batch_size=32, num_solutions = None) -> tuple[torch.Tensor, torch.Tensor]:
     assert input_ids.shape[0] == 1, "Batch size must be 1"
     
     if num_solutions is None:
         num_solutions = beam_size
+    vanilla = multi_choices_parser is None
+    if vanilla:
+        multi_choices_parser = AcceptEverythingParser(model.config.vocab_size)
 
     parsers_unfinished = [multi_choices_parser]
     scores_finished = torch.tensor([], dtype=torch.float)
@@ -73,9 +90,10 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
         logprobs_filtered = apply_mask_tokens(logprobs, parsers_tokens)
         if len(logprobs_filtered):
             topk = torch.topk(logprobs_filtered, beam_size, dim=-1) # shape (batch_size, beam_size)
-            topk_global = topk.values.flatten().topk(beam_size)
+            values = topk.values + scores_unfinished.unsqueeze(-1)
+            topk_global = values.flatten().topk(beam_size)
             best_tokens_row = topk_global.indices // beam_size
-            best_tokens, best_tokens_logprobs = topk.indices[best_tokens_row, topk_global.indices % beam_size], topk_global.values
+            best_tokens, best_tokens_logprobs = topk.indices[best_tokens_row, topk_global.indices % beam_size], topk.values[best_tokens_row, topk_global.indices % beam_size]
             notinf = ~best_tokens_logprobs.isinf()
             best_tokens, best_tokens_row, best_tokens_logprobs = best_tokens[notinf], best_tokens_row[notinf], best_tokens_logprobs[notinf]
         else:
@@ -104,9 +122,11 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
         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)
+
+    # Special case of vanilla beam search where all answers are valid
+    if vanilla:
+        order = scores_unfinished.argsort(descending=True)
+        scores_finished = scores_unfinished[order][:num_solutions]
+        solutions_finished = solutions_unfinished[order][:num_solutions]
     
     return scores_finished, solutions_finished
-        
-
-        
-        

divergent_beamsearch-0.1.1/tests/test_beamsearch.py

@@ -0,0 +1,118 @@
+import numpy as np
+import pytest
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from multi_choices_parser import MultiChoicesParser
+from divergent_beamsearch.algorithm import divergent_beamsearch, log1mexp
+from multi_choices_parser import MultiChoicesParser
+
+@pytest.fixture
+def model_and_tokenizer():
+    model = GPT2LMHeadModel.from_pretrained("gpt2")
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    return model, tokenizer
+
+def test_divergent_beamsearch(model_and_tokenizer):
+    model, tokenizer = model_and_tokenizer
+    prompt = "The capital of France is"
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    beam_size = 5
+    max_length = 10
+    pad_token_id = tokenizer.eos_token_id
+
+    possible_answers = [' Paris', ' Paris Hilton']
+    tokenized_answers = tokenizer(possible_answers).input_ids
+    multi_choices_parser = MultiChoicesParser([tokenized_answers])
+
+    logprob_paris = model(input_ids).logits.log_softmax(dim=-1)[0, -1, tokenized_answers[0][0]]
+    logprob_hilton = model(torch.cat([input_ids, torch.tensor(tokenized_answers[1][0]).view(1,1)], dim=-1)).logits.log_softmax(dim=-1)[0, -1, tokenized_answers[1][1]]
+    logprob_paris_hilton = logprob_paris + logprob_hilton
+
+    scores, solutions = divergent_beamsearch(
+        input_ids=input_ids,
+        model=model,
+        beam_size=beam_size,
+        max_length=max_length,
+        multi_choices_parser=multi_choices_parser,
+        pad_token_id=pad_token_id,
+        num_solutions=10
+    )
+    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"
+
+def test_vanilla_beamsearch(model_and_tokenizer):
+    # Verify that divergent beam search where all answers are valid is equivalent to vanilla beam search 
+    # Results of beam search were compared with huggingface implementation (https://huggingface.co/spaces/m-ric/beam_search_visualizer)
+    model, tok = model_and_tokenizer
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.eval()
+    prompt = "The capital of France is"
+    input_ids = tok(prompt, return_tensors="pt").input_ids.to(device)
+    scores, sequences = divergent_beamsearch(
+        input_ids, model, beam_size=3, max_length=1, pad_token_id=tok.eos_token_id, num_solutions=3, multi_choices_parser=None
+    )
+    sequences = [tok.decode(s) for s in sequences]
+    assert sequences == [" the", " now", " a"]
+    assert np.isclose(
+        scores.cpu().numpy(), np.array([-2.4699, -3.0377, -3.0756]), atol=0.0001
+    ).all()
+
+    scores, sequences = divergent_beamsearch(
+        input_ids, model, beam_size=3, max_length=2, pad_token_id=tok.eos_token_id, num_solutions=3, multi_choices_parser=None
+    )
+    sequences = [tok.decode(s) for s in sequences]
+    assert sequences == [" the capital", " now home", " now the"]
+    assert np.isclose(
+        scores.cpu().numpy(), np.array([-4.2437, -5.3013, -5.3408]), atol=0.0001
+    ).all()
+
+    scores, sequences = divergent_beamsearch(
+        input_ids, model, beam_size=3, max_length=3, pad_token_id=tok.eos_token_id, num_solutions=3, multi_choices_parser=None
+    )
+    sequences = [tok.decode(s) for s in sequences]
+    assert sequences == [" the capital of", " now home to", " now the capital"]
+    assert np.isclose(
+        scores.cpu().numpy(), np.array([-4.3194, -5.3057, -7.7173]), atol=0.0001
+    ).all()
+
+    scores, sequences = divergent_beamsearch(
+        input_ids, model, beam_size=3, max_length=4, pad_token_id=tok.eos_token_id, num_solutions=3, multi_choices_parser=None
+    )
+    sequences = [tok.decode(s) for s in sequences]
+    assert sequences == [
+        " the capital of the",
+        " the capital of France",
+        " the capital of a",
+    ]
+    assert np.isclose(
+        scores.cpu().numpy(), np.array([-5.5825, -5.9150, -7.1716]), atol=0.0001
+    ).all()
+
+    scores, sequences = divergent_beamsearch(
+        input_ids, model, beam_size=3, max_length=5, pad_token_id=tok.eos_token_id, num_solutions=3, multi_choices_parser=None
+    )
+    sequences = [tok.decode(s) for s in sequences]
+    assert sequences == [
+        " the capital of France,",
+        " the capital of France.",
+        " the capital of the French",
+    ]
+    assert np.isclose(
+        scores.cpu().numpy(), np.array([-6.9453, -7.1549, -7.5727]), atol=0.0001
+    ).all()
+
+
+    scores, sequences = divergent_beamsearch(
+        input_ids, model, beam_size=3, max_length=6, pad_token_id=tok.eos_token_id, num_solutions=3, multi_choices_parser=None
+    )
+    sequences = [tok.decode(s) for s in sequences]
+    assert sequences == [
+        " the capital of France, and",
+        " the capital of the French Republic",
+        " the capital of France. It",
+    ]
+    assert np.isclose(
+        scores.cpu().numpy(), np.array([-8.1361, -8.7745, -9.1053]), atol=0.0001
+    ).all()

{divergent_beamsearch-0.1.0 → divergent_beamsearch-0.1.1}/uv.lock

@@ -73,7 +73,7 @@ wheels = [
 
 [[package]]
 name = "divergent-beamsearch"
-version = "0.1.0"
+version = "0.1.1"
 source = { editable = "." }
 dependencies = [
     { name = "multi-choices-parser" },

divergent_beamsearch-0.1.0/tests/test_beamsearch.py

@@ -1,41 +0,0 @@
-import pytest
-import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-from multi_choices_parser import MultiChoicesParser
-from divergent_beamsearch.algorithm import divergent_beamsearch, log1mexp
-
-@pytest.fixture
-def model_and_tokenizer():
-    model = GPT2LMHeadModel.from_pretrained("gpt2")
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    return model, tokenizer
-
-def test_divergent_beamsearch(model_and_tokenizer):
-    model, tokenizer = model_and_tokenizer
-    prompt = "The capital of France is"
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    beam_size = 5
-    max_length = 10
-    pad_token_id = tokenizer.eos_token_id
-
-    possible_answers = [' Paris', ' Paris Hilton']
-    tokenized_answers = tokenizer(possible_answers).input_ids
-    multi_choices_parser = MultiChoicesParser([tokenized_answers])
-
-    logprob_paris = model(input_ids).logits.log_softmax(dim=-1)[0, -1, tokenized_answers[0][0]]
-    logprob_hilton = model(torch.cat([input_ids, torch.tensor(tokenized_answers[1][0]).view(1,1)], dim=-1)).logits.log_softmax(dim=-1)[0, -1, tokenized_answers[1][1]]
-    logprob_paris_hilton = logprob_paris + logprob_hilton
-
-    scores, solutions = divergent_beamsearch(
-        input_ids=input_ids,
-        model=model,
-        beam_size=beam_size,
-        max_length=max_length,
-        multi_choices_parser=multi_choices_parser,
-        pad_token_id=pad_token_id,
-        num_solutions=10
-    )
-    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"