divergent-beamsearch 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

divergent_beamsearch/algorithm.py

@@ -1,7 +1,7 @@
 import math
 import torch
 from transformers import GPT2LMHeadModel
-from multi_choices_parser import MultiChoicesParser, end_symb
+from multi_choices_parser import DEFAULT_END_SYMB
 
 
 class Parser:
@@ -19,10 +19,10 @@ def get_parsers_tokens(parsers : list[Parser], end_symb) -> tuple[list, list[int
     can_end = []
     for parser in parsers:
         tokens = list(parser.next())
-        if end_symb in tokens:
-            can_end.append(True)
+        try:
             tokens.remove(end_symb)
-        else:
+            can_end.append(True)
+        except ValueError:         
             can_end.append(False)
         parsers_tokens.append(tokens)
     return parsers_tokens, can_end
@@ -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, end_symb=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) -> 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()
@@ -140,6 +140,7 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
             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
     if vanilla:
         order = scores_unfinished.argsort(descending=True)
         scores_finished = scores_unfinished[order][:num_solutions]
@@ -154,7 +155,9 @@ 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, parsers : Parser | list[Parser] | None, batch_size=32, start : int | torch.IntTensor = None) -> torch.FloatTensor:
+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:
     if start is None:
         start = 0
     if isinstance(start, int):

{divergent_beamsearch-0.1.5.dist-info → divergent_beamsearch-0.1.6.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.1.5
+Version: 0.1.6
 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.57
+Requires-Dist: multi-choices-parser>=0.9.61
 Requires-Dist: torch>=2.0.0
 Requires-Dist: transformers>=4.47.1
 Description-Content-Type: text/markdown

divergent_beamsearch-0.1.6.dist-info/RECORD

@@ -0,0 +1,6 @@
+divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
+divergent_beamsearch/algorithm.py,sha256=73BsS5IU1_4Aj11LdQebKofpCO-Mo8BXtDfx-AEYxXA,8835
+divergent_beamsearch-0.1.6.dist-info/METADATA,sha256=cm-VyQfzk9sklvIkFXgEfI4A4ktWddIO5CvtaL7Vkng,2826
+divergent_beamsearch-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+divergent_beamsearch-0.1.6.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
+divergent_beamsearch-0.1.6.dist-info/RECORD,,

divergent_beamsearch-0.1.5.dist-info/RECORD

@@ -1,6 +0,0 @@
-divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
-divergent_beamsearch/algorithm.py,sha256=b44kA2_M0HfOSC6LYQmu7W_JnGAl1u9Sz_91jMjvWg0,8688
-divergent_beamsearch-0.1.5.dist-info/METADATA,sha256=sRitbfgDp8YqLXZzZmZw7Nd1pJWv8T4HJoSqjhO4ztI,2826
-divergent_beamsearch-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-divergent_beamsearch-0.1.5.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
-divergent_beamsearch-0.1.5.dist-info/RECORD,,