Test Flashlight Text 12/22/22

Author: joecummings

test/torchtext_unittest/prototype/test_generate.py

@@ -16,10 +16,10 @@ def setUp(self) -> None:
         self.inputs = self.transform(
             [
                 "summarize: studies have shown that owning a dog is good for you",
-                "translate English to German: That is good.",
-                "cola sentence: The course is jumping well.",
-                "stsb sentence1: The rhino grazed on the grass. sentence2: A rhino is grazing in a field.",
-                "summarize: state authorities dispatched emergency crews tuesday to survey the damage after an onslaught of severe weather in mississippi...",
+                # "translate English to German: That is good.",
+                # "cola sentence: The course is jumping well.",
+                # "stsb sentence1: The rhino grazed on the grass. sentence2: A rhino is grazing in a field.",
+                # "summarize: state authorities dispatched emergency crews tuesday to survey the damage after an onslaught of severe weather in mississippi...",
             ]
         )
         torch.manual_seed(0)
@@ -50,4 +50,20 @@ def test_generate_errors_with_incorrect_beams(self) -> None:
     def test_warns_when_no_max_len_provided(self, mock) -> None:
         generation_model = GenerationUtil(self.model)
         generation_model.generate(self.inputs)
-        mock.assert_called_with("`max_len` was not specified. Defaulting to 100 tokens.")
+        mock.assert_called_with("`max_len` was not specified. Defaulting to 256 tokens.")
+    
+    def test_beam_search(self) -> None:
+        generation_model = GenerationUtil(self.model)
+
+        tokens = generation_model.generate(self.inputs, num_beams=3, max_len=30)
+
+        import pdb
+
+        pdb.set_trace()
+
+        generated_text = self.transform.decode(tokens.tolist())
+
+        import pdb
+        pdb.set_trace()
+
+

torchtext/prototype/generate.py

@@ -1,14 +1,25 @@
-from typing import Optional
+from dataclasses import dataclass
+from typing import List, Optional
 
 import torch
 import torch.nn.functional as F
 from torch import nn
+from flashlight.lib.text.decoder import LexiconFreeSeq2SeqDecoder, LexiconFreeSeq2SeqDecoderOptions, ZeroLM, create_emitting_model_state, get_obj_from_emitting_model_state
 
 import logging
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class Seq2SeqModelState(object):
+
+    timestep: int
+    hidden_states: List[torch.Tensor]
+    sequence: torch.Tensor
+    lm_scores: torch.Tensor
+
+
 class GenerationUtil:
     """Wrapper to provide generation utils for encoder/decoder models and decoder models.
 
@@ -92,8 +103,90 @@ def greedy_search(
 
         return input_ids
 
-    def beam_search(self, input_ids: torch.Tensor, num_beams: int, max_len: Optional[int]) -> torch.Tensor:
-        raise NotImplementedError()
+    def beam_search(
+        self, 
+        input_ids: torch.Tensor,
+        num_beams: int,
+        max_len: int,
+        vocab_size: int,
+        eos_idx: int = 1,
+        **model_kwargs
+    ) -> torch.Tensor:
+        
+        # Is this right?
+        T = max_len
+        N = vocab_size
+
+        def update_func(emissions_ptr, N, T, prev_step_token_idxs, prev_step_model_states, timestep):
+            # `emissions_ptr` should always be the same (from encoder output)
+            # N is not needed
+            # T is not needed
+            
+            if timestep == 0:
+                prev_step_token_idxs = input_ids
+                prev_step_model_states = [
+                    create_emitting_model_state(
+                        Seq2SeqModelState(
+                            timestep=0,
+                            hidden_states=None,
+                            sequence=input_ids,
+                            lm_scores=None
+                        )
+                    )
+                ]
+
+            model_inputs = self.model.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            if self.is_huggingface_model:
+                model_inputs["return_dict"] = True
+                model_inputs["output_hidden_states"] = True
+
+            outputs = self.model(**model_inputs)
+            output_key = "logits" if self.is_huggingface_model else "decoder_output"
+            lm_scores = outputs[output_key]
+
+            model_states = []
+            for idx, model_state_ptr in zip(prev_step_token_idxs, prev_step_model_states):
+                model_state = get_obj_from_emitting_model_state(model_state_ptr)
+                model_states.append(
+                    create_emitting_model_state(
+                        Seq2SeqModelState(
+                            timestep=timestep,
+                            hidden_states=outputs["decoder_hidden_states"],
+                            sequence=torch.cat([model_state.sequence[:, -1], idx], dim=-1),
+                            lm_scores=lm_scores
+                        )
+                    )
+                )
+
+            import pdb
+            pdb.set_trace()
+
+            out_probs = lm_scores[0][0].tolist() * len(prev_step_token_idxs)
+            return out_probs, model_states
+
+        options = LexiconFreeSeq2SeqDecoderOptions(
+            beam_size=num_beams,
+            beam_size_token=self.model.config.vocab_size,
+            beam_threshold=1000,
+            lm_weight=0.0,
+            eos_score=0.0,
+            log_add=True,
+        )
+
+        decoder = LexiconFreeSeq2SeqDecoder(
+            options=options,
+            lm=ZeroLM(),
+            eos_idx=eos_idx,
+            update_func=update_func,
+            max_output_length=max_len
+        )
+
+        emissions = model_kwargs["encoder_outputs"].get("encoder_output")
+
+        decoder.decode_step(emissions.data_ptr(), T, N)
+        hyps = decoder.get_all_final_hypothesis()
+
+        return hyps
 
     def generate(
         self,
@@ -135,6 +228,9 @@ def generate(
         if num_beams == 1 or num_beams is None:
             return self.greedy_search(inputs, max_len, eos_idx, pad_idx=pad_idx, **model_kwargs)
         elif num_beams > 1:
-            return self.beam_search(inputs, num_beams, max_len)
+            if torch.has_cuda:
+                logger.warning("No CUDA parellelization available through CUDA yet.")
+            # Implement some sort of multiprocessing here
+            return self.beam_search(inputs, num_beams, vocab_size=self.model.config.vocab_size, max_len=max_len, **model_kwargs)
         else:
             raise ValueError("`num_beams` must be >= 1.")