divergent-beamsearch 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

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.dist-info → divergent_beamsearch-0.2.2.dist-info}/METADATA

@@ -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.2.dist-info/RECORD

@@ -0,0 +1,6 @@
+divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
+divergent_beamsearch/algorithm.py,sha256=sfQSaOcnuN_Pyp69hGZbX66Rmg7V4aQqyUfT0Clltvs,10204
+divergent_beamsearch-0.2.2.dist-info/METADATA,sha256=l9Wsb45vwcLSNMcsiTuGay5rq9o4TD_Pt0D7ZnipTTg,2790
+divergent_beamsearch-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+divergent_beamsearch-0.2.2.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
+divergent_beamsearch-0.2.2.dist-info/RECORD,,

divergent_beamsearch-0.2.0.dist-info/RECORD

@@ -1,6 +0,0 @@
-divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
-divergent_beamsearch/algorithm.py,sha256=lx27rXddHiyzisINgWI5MuatRLIU2ObnZhtCvojbGJ8,9917
-divergent_beamsearch-0.2.0.dist-info/METADATA,sha256=u4-bH-9qa_yLJPemATIemwIavOCucF7CCv0kyJV6_Qg,2826
-divergent_beamsearch-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-divergent_beamsearch-0.2.0.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
-divergent_beamsearch-0.2.0.dist-info/RECORD,,