[BART] Do not add start/end tokens multiple times

parlai/agents/bart/bart.py

@@ -158,13 +158,19 @@ def _set_text_vec(
         if 'text' not in obs or 'text_vec' not in obs:
             return obs
         vec = obs['text_vec']
-        if truncate is not None:
-            vec = torch.LongTensor(  # type: ignore
-                self._check_truncate(obs['text_vec'], truncate - 2, True)
+
+        # add start/end tokens
+        if 'added_start_end_tokens' not in obs:
+            if truncate is not None:
+                vec = torch.LongTensor(  # type: ignore
+                    self._check_truncate(obs['text_vec'], truncate - 2, True)
+                )
+            obs.force_set(
+                'text_vec',
+                self._add_start_end_tokens(vec, add_start=True, add_end=True),
             )
-        obs.force_set(
-            'text_vec', self._add_start_end_tokens(vec, add_start=True, add_end=True)
-        )
+            obs['added_start_end_tokens'] = True
+
         return obs
 
     def _get_initial_decoder_input(

tests/nightly/gpu/test_bart.py

@@ -41,6 +41,27 @@ def test_bart_gen(self):
 
         self.assertEqual(act['text'], text)
 
+    def test_bart_cache_text_vec(self):
+        """
+        Test BART text vec caching
+        """
+        opt = ParlaiParser(True, True).parse_args(['--model', 'bart'])
+        bart = create_agent(opt)
+
+        # obs 1
+        text = "Don't have a cow, Man!"
+        in_obs = {'text': text, 'episode_done': True}
+        out_obs = bart.observe(in_obs)
+        cached_text_vec = out_obs['text_vec']
+        _ = bart.act()
+
+        # obs 2
+        in_obs = {'text_vec': cached_text_vec, 'episode_done': True}
+        out_obs = bart.observe(in_obs)
+        cached_text_vec_2 = out_obs['text_vec']
+
+        self.assertEqual(cached_text_vec.tolist(), cached_text_vec_2.tolist())
+
     @testing_utils.retry(ntries=3, log_retry=True)
     def test_bart_ft(self):
         """