divergent-beamsearch 0.1.4__tar.gz → 0.1.5__tar.gz

{divergent_beamsearch-0.1.4 → divergent_beamsearch-0.1.5}/PKG-INFO

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

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

@@ -14,7 +14,7 @@ class Parser:
     def copy(self):
         raise NotImplementedError
 
-def get_parsers_tokens(parsers : list[Parser]) -> tuple[list, list[int]]:
+def get_parsers_tokens(parsers : list[Parser], end_symb) -> tuple[list, list[int]]:
     parsers_tokens = []
     can_end = []
     for parser in parsers:
@@ -77,7 +77,7 @@ class AcceptEverythingParser(Parser):
         return self
 
 @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) -> 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=end_symb) -> 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()
@@ -101,7 +101,7 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
         if len(input_ids_unfinished) == 0:
             break
         pred = batched_inference_logits(model, input_ids_unfinished.to(device), batch_size=batch_size)[:, -1].cpu()
-        parsers_tokens, can_end = get_parsers_tokens(parsers_unfinished)
+        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)
         if len(logprobs_filtered):