chkpt

Author: joecummings

test/torchtext_unittest/prototype/test_generate.py

@@ -55,7 +55,7 @@ def test_warns_when_no_max_len_provided(self, mock) -> None:
     def test_beam_search(self) -> None:
         generation_model = GenerationUtil(self.model)
 
-        tokens = generation_model.generate(self.inputs, num_beams=3, max_len=30)
+        tokens = generation_model.generate(self.inputs, num_beams=3, max_len=100)
 
         generated_text = self.transform.decode(tokens.tolist())
 

torchtext/prototype/generate.py

@@ -83,6 +83,7 @@ def greedy_search(
             decoder_output = outputs[output_key]
 
             # Calculate probabilities and take the most likely next token
+            # Why do we take the last token instead of a mean_pooling across all of them?
             probs = F.log_softmax(decoder_output[:, -1], dim=-1)
             _, next_tokens = torch.topk(probs, 1)
 
@@ -129,38 +130,48 @@ def update_func(emissions_ptr, N, T, prev_step_token_idxs, prev_step_model_state
                         Seq2SeqModelState(
                             timestep=0,
                             hidden_states=None,
-                            sequence=input_ids[:, -1],
+                            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 = []
+            out_probs, model_states = [], []
             for idx, model_state_ptr in zip(prev_step_token_idxs, prev_step_model_states):
                 if isinstance(idx, int):
-                    idx = torch.Tensor([idx])
-                model_state = get_obj_from_emitting_model_state(model_state_ptr)
+                    idx = torch.Tensor([idx]).to(torch.long)
+
+                # Get previous model state
+                prev_model_state = get_obj_from_emitting_model_state(model_state_ptr)
+                
+                # Create new decoder token ids
+                new_input_ids = torch.cat([prev_model_state.sequence[:, -1], idx], dim=-1)
+                
+                # Forward pass
+                model_inputs = self.model.prepare_inputs_for_generation(new_input_ids.unsqueeze(dim=0), **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]
+
+                # Keep track of probabilities over vocab for this pairing
+                out_probs.append(torch.squeeze(lm_scores[:, -1]).tolist())
+                
+                # Keep track of sequence and decoder hidden states
                 model_states.append(
                     create_emitting_model_state(
                         Seq2SeqModelState(
                             timestep=timestep,
                             hidden_states=outputs["decoder_hidden_states"],
-                            sequence=torch.cat([model_state.sequence, idx], dim=-1),
+                            sequence=new_input_ids.unsqueeze(dim=0),
                             lm_scores=lm_scores
                         )
                     )
                 )
 
-            out_probs = lm_scores[0].tolist() * len(prev_step_token_idxs)
             return out_probs, model_states
 
         options = LexiconFreeSeq2SeqDecoderOptions(
@@ -188,7 +199,13 @@ def update_func(emissions_ptr, N, T, prev_step_token_idxs, prev_step_model_state
         token_scores = [(hyp.tokens, hyp.score) for hyp in hyps]
         max_tokens = max(token_scores, key=lambda x: x[1])
 
-        return torch.Tensor(max_tokens[0]).to(torch.int)
+        filtered = list(filter(lambda x: x != -1, max_tokens[0]))
+        final_tokens = [0] + filtered
+        
+        import pdb
+        pdb.set_trace()
+
+        return torch.Tensor(final_tokens).to(torch.long)
 
     def generate(
         self,