divergent-beamsearch 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

divergent_beamsearch/algorithm.py

@@ -1,4 +1,5 @@
 import math
+import multi_choices_parser
 import torch
 try:
     from transformers import GPT2LMHeadModel
@@ -74,6 +75,7 @@ class AcceptEverythingParser(Parser):
     def __init__(self, vocab_size : int):
         self.vocab_size = vocab_size
         self.tokens = tuple(range(vocab_size))
+        self.finished = False
 
     def step(self, token):
         pass
@@ -163,9 +165,11 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : "GPT2LMHeadModel", be
         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
@@ -187,6 +191,7 @@ def divergent_logprob(input_ids : torch.Tensor, attention_mask : torch.Tensor |
                       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])
@@ -225,8 +230,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.1.dist-info → divergent_beamsearch-0.2.3.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.2.1
+Version: 0.2.3
 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.72
+Requires-Dist: multi-choices-parser>=0.10.0
 Requires-Dist: torch>=2.0.0
 Description-Content-Type: text/markdown
 

divergent_beamsearch-0.2.3.dist-info/RECORD

@@ -0,0 +1,6 @@
+divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
+divergent_beamsearch/algorithm.py,sha256=n41WYmcT9UwuCBHbdKm3-SgO3NgJBenTlueVxpF1dV8,10234
+divergent_beamsearch-0.2.3.dist-info/METADATA,sha256=o33zn3fuy96QWws-0ALVXo9euVKug5vXbtWDZ4eF6t8,2790
+divergent_beamsearch-0.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+divergent_beamsearch-0.2.3.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
+divergent_beamsearch-0.2.3.dist-info/RECORD,,

divergent_beamsearch-0.2.1.dist-info/RECORD

@@ -1,6 +0,0 @@
-divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
-divergent_beamsearch/algorithm.py,sha256=GKFwi6aKNmJRu9SR6X96JT93SbOpy84fxyKJ5Pq5vQs,9961
-divergent_beamsearch-0.2.1.dist-info/METADATA,sha256=0JAVae-tlHYFQkaEqBOE9ZDtExKsS-gpFFFb9oNTRdg,2790
-divergent_beamsearch-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-divergent_beamsearch-0.2.1.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
-divergent_beamsearch-0.2.1.dist-info/RECORD,,