ckkpt

Author: joecummings

test/torchtext_unittest/prototype/test_generate.py

@@ -57,10 +57,6 @@ def test_beam_search(self) -> None:
 
         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

torchtext/prototype/generate.py

@@ -129,7 +129,7 @@ def update_func(emissions_ptr, N, T, prev_step_token_idxs, prev_step_model_state
                         Seq2SeqModelState(
                             timestep=0,
                             hidden_states=None,
-                            sequence=input_ids,
+                            sequence=input_ids[:, -1],
                             lm_scores=None
                         )
                     )
@@ -146,30 +146,29 @@ def update_func(emissions_ptr, N, T, prev_step_token_idxs, prev_step_model_state
 
             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)
                 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),
+                            sequence=torch.cat([model_state.sequence, idx], dim=-1),
                             lm_scores=lm_scores
                         )
                     )
                 )
 
-            import pdb
-            pdb.set_trace()
-
-            out_probs = lm_scores[0][0].tolist() * len(prev_step_token_idxs)
+            out_probs = lm_scores[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,
+            beam_threshold=50,
             lm_weight=0.0,
-            eos_score=0.0,
+            eos_score=1.0,
             log_add=True,
         )
 
@@ -186,7 +185,10 @@ def update_func(emissions_ptr, N, T, prev_step_token_idxs, prev_step_model_state
         decoder.decode_step(emissions.data_ptr(), T, N)
         hyps = decoder.get_all_final_hypothesis()
 
-        return hyps
+        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)
 
     def generate(
         self,