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

divergent_beamsearch/algorithm.py

@@ -35,12 +35,17 @@ 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, attention_mask : torch.Tensor | None = None, batch_size : int = 32) -> 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 = []
     if attention_mask is None:
         attention_mask = torch.ones_like(input_ids)
     for i in range(0, input_ids.shape[0], batch_size):
-        logits.append(model(input_ids[i:i+batch_size], attention_mask=attention_mask[i:i+batch_size]).logits)
+        l = model(input_ids[i:i+batch_size], attention_mask=attention_mask[i:i+batch_size]).logits
+        if to_cpu:
+            l = l.cpu()
+        logits.append(l)
     return torch.cat(logits, dim=0)
 
 def select_mask(source : list, mask : list[bool]) -> list:
@@ -91,7 +96,9 @@ 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, 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]:
+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"
     device = input_ids.device
     input_ids = input_ids.cpu()
@@ -114,7 +121,7 @@ def divergent_beamsearch(input_ids : torch.Tensor, model : GPT2LMHeadModel, beam
     for _ in range(max_length):
         if len(input_ids_unfinished) == 0:
             break
-        pred = batched_inference_logits(model, input_ids_unfinished.to(device), batch_size=batch_size)[:, -1].cpu()
+        pred = batched_inference_logits(model, input_ids_unfinished.to(device), batch_size=batch_size, to_cpu=optimize_gpu_mem)[:, -1].cpu()
         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)
@@ -175,19 +182,20 @@ def set_slice_row(x : torch.Tensor, slices : torch.IntTensor, value) -> torch.Te
 @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, end_symb=DEFAULT_END_SYMB) -> torch.FloatTensor:
+                      start : int | torch.IntTensor = None, end_symb=DEFAULT_END_SYMB, optimize_gpu_mem=True) -> torch.FloatTensor:
     if start is None:
-        start = 0
+        start = 1
     if isinstance(start, int):
         start = torch.tensor([start]*input_ids.shape[0])
     assert start.shape[0] == input_ids.shape[0]
+    assert (start > 0).all()
     # -1 because next token offset
     start = start - 1
 
     if attention_mask is None:
         attention_mask = torch.ones_like(input_ids)
 
-    logits = batched_inference_logits(model, input_ids, attention_mask, batch_size).cpu()
+    logits = batched_inference_logits(model, input_ids, attention_mask, batch_size, to_cpu=optimize_gpu_mem).cpu()
     input_ids = input_ids.cpu()
     attention_mask = attention_mask.cpu()
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: divergent-beamsearch
-Version: 0.1.8
+Version: 0.2.0
 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.2.0.dist-info/RECORD

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

divergent_beamsearch-0.1.8.dist-info/RECORD

@@ -1,6 +0,0 @@
-divergent_beamsearch/__init__.py,sha256=qrpVRoT3d-q1N9fJnzHI2X13e71LDY4-6eLOQ_gwCqQ,62
-divergent_beamsearch/algorithm.py,sha256=rywmvaIoo66aksaNdCXOPfqtd8WnCazVqYoxySi6G9s,9610
-divergent_beamsearch-0.1.8.dist-info/METADATA,sha256=iZjtT-uUwN1X2EfFzPI5_ermjIMu9Myz3d4H8FWR4nw,2826
-divergent_beamsearch-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-divergent_beamsearch-0.1.8.dist-info/licenses/LICENCE,sha256=gnISbTzmuQC7NwJaGOdjoq26QYgSuKndq5q2JykifKw,1075
-divergent_beamsearch-0.1.8.dist-info/RECORD,,