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

divergent_beamsearch/algorithm.py

@@ -1,6 +1,9 @@
 import math
 import torch
-from transformers import GPT2LMHeadModel
+try:
+    from transformers import GPT2LMHeadModel
+except ImportError:
+    pass
 from multi_choices_parser import DEFAULT_END_SYMB
 
 
@@ -35,7 +38,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 +99,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"
@@ -180,7 +183,7 @@ 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:

{divergent_beamsearch-0.2.0.dist-info → divergent_beamsearch-0.2.1.dist-info}/METADATA

@@ -1,12 +1,11 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.2.0
+Version: 0.2.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
-Requires-Dist: multi-choices-parser>=0.9.61
+Requires-Dist: multi-choices-parser>=0.9.72
 Requires-Dist: torch>=2.0.0
-Requires-Dist: transformers>=4.47.1
 Description-Content-Type: text/markdown
 
 # Divergent Beam Search

divergent_beamsearch-0.2.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+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,,

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,,