We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Wizard of Wikipedia] Where is get_dialog_history() defined?
https://github.com/facebookresearch/ParlAI/blob/a43f2880719c5a048fdf3d0aa5d5b25eeb9a1a41/projects/wizard_of_wikipedia/generator/agents.py#L78

@stephenroller I'm guessing this got missed in TorchAgent. Could you please fix it?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `projects/wizard_of_wikipedia/generator/agents.py`
Content:
```
1 #!/usr/bin/env python
2 
3 # Copyright (c) Facebook, Inc. and its affiliates.
4 # This source code is licensed under the MIT license found in the
5 # LICENSE file in the root directory of this source tree.
6 
7 """
8 Agents for handling the generation aspects of Wizard.
9 """
10 from itertools import chain
11 from functools import lru_cache
12 
13 import torch as th
14 import numpy as np
15 
16 from parlai.core.utils import padded_tensor, round_sigfigs
17 
18 from parlai.agents.transformer.transformer import TransformerGeneratorAgent
19 
20 from .modules import EndToEndModel
21 from parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE
22 
23 TOKEN_DIALOG = '__dialog__'
24 
25 
26 DEFAULT_OPTS = {
27     "learningrate": 5e-4,
28     "optimizer": "adam",
29     "lr_scheduler": "invsqrt",
30     "warmup_updates": 5000,
31     "clip_norm": 0.1,
32     "ffn_size": 512,
33     "embedding_size": 256,
34     "n_heads": 2,
35     "dropout": 0.2,
36     "n_layers": 5,
37     "betas": "0.9,0.98",
38     "truncate": 128,
39     "dict_textfields": "text,labels,chosen_topic,checked_sentence,knowledge,title",
40 }
41 
42 
43 class _GenericWizardAgent(TransformerGeneratorAgent):
44     @classmethod
45     def add_cmdline_args(cls, argparser):
46         argparser.set_defaults(**DEFAULT_OPTS)
47         super(_GenericWizardAgent, cls).add_cmdline_args(argparser)
48 
49     def batchify(self, obs_batch):
50         batch = super().batchify(obs_batch)
51         reordered_observations = [obs_batch[i] for i in batch.valid_indices]
52 
53         checked_sentences = []
54         for obs in reordered_observations:
55             checked_sentence = '{} {} {}'.format(
56                 obs.get('title', ''), TOKEN_KNOWLEDGE, obs.get('checked_sentence', '')
57             )
58             checked_sentences.append(checked_sentence)
59 
60         batch['checked_sentence'] = checked_sentences
61         return batch
62 
63 
64 class TwoStageAgent(_GenericWizardAgent):
65     def __init__(self, opt, shared):
66         super().__init__(opt, shared)
67         if shared is not None:
68             # make sure the dialogue token appears
69             self.dict[TOKEN_DIALOG] = 9999999
70 
71     def observe(self, obs):
72         if 'text' not in obs:
73             return obs
74 
75         # TODO: resolve this with #1421
76         # get the dialog stuff
77         reply = self.last_reply()
78         self.observation = self.get_dialog_history(obs, reply=reply)
79         # we need to store the old text so that we can restore it
80         oldtext = obs['text']
81 
82         # now we want to force prepend the knowledge stuff
83         fields = []
84         if 'chosen_topic' in obs:
85             fields += [obs['title']]
86         if 'checked_sentence' in obs:
87             fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]
88         if obs['text'] != '':
89             fields += [TOKEN_DIALOG, obs['text']]
90         obs['text'] = ' '.join(fields)
91 
92         # now vectorize with the extra knowledge. It'll all get stored in the
93         # text_vec operation, etc
94         self.vectorize(
95             obs,
96             text_truncate=self.text_truncate,
97             label_truncate=self.label_truncate
98         )
99 
100         # finally we need to return the old text to the way it was
101         obs['text'] = oldtext
102         assert obs is self.observation
103 
104         return obs
105 
106 
107 class EndToEndAgent(_GenericWizardAgent):
108     def __init__(self, opt, shared=None):
109         super().__init__(opt, shared)
110         self._vectorize_text = lru_cache(int(2 ** 20))(self._vectorize_text)
111 
112         # knowledge truncate defaults to the same as --truncate
113         self.knowledge_truncate = opt.get('knowledge_truncate')
114         if not self.knowledge_truncate:
115             self.knowledge_truncate = opt['truncate']
116         self.max_knowledge = opt.get('max_knowledge')
117         self.knowledge_alpha = opt['knowledge_alpha']
118 
119     def _dummy_batch(self, bsz, maxlen):
120         batch = super()._dummy_batch(bsz, maxlen)
121         batch['know_vec'] = th.zeros(bsz, 2, 2).long().cuda()
122         batch['ck_mask'] = th.ones(bsz, 2).byte().cuda()
123         batch['cs_ids'] = th.zeros(bsz).long().cuda()
124         batch['use_cs_ids'] = True
125         return batch
126 
127     def compute_loss(self, batch, return_output=False):
128         # first compute our regular forced decoding loss
129         token_loss, model_output = super().compute_loss(batch, return_output=True)
130         notnull = batch.label_vec.ne(self.NULL_IDX)
131         num_tokens = notnull.long().sum().item()
132 
133         encoder_states = model_output[2]
134         ctx_know_attn = encoder_states[2]
135 
136         if self.knowledge_alpha == 0.0:
137             loss = token_loss
138         else:
139             _, know_pred = ctx_know_attn.max(1)
140             know_acc = (know_pred == batch.cs_ids).float().sum().item()
141             know_chance = batch.ck_mask.sum(1).float().reciprocal().sum().item()
142             self.metrics['know_chance'] += know_chance
143             self.metrics['bsz'] += batch.text_vec.size(0)
144             self.metrics['know_acc'] += know_acc
145             know_loss = th.nn.functional.cross_entropy(
146                 ctx_know_attn,
147                 batch.cs_ids,
148                 reduction='mean',
149             )
150             self.metrics['know_loss'] += know_loss.item() * batch.text_vec.size(0)
151             # in the original paper the loss was scaled by num_tokens for both
152             # know_loss and token_loss
153             know_loss /= num_tokens
154             loss = (
155                 (1 - self.knowledge_alpha) * token_loss +
156                 self.knowledge_alpha * know_loss
157             )
158         if return_output:
159             return (loss, model_output)
160         else:
161             return loss
162 
163     def reset_metrics(self):
164         super().reset_metrics()
165         self.metrics['bsz'] = 0.0
166         self.metrics['know_acc'] = 0.0
167         self.metrics['know_loss'] = 0.0
168         self.metrics['know_chance'] = 0.0
169 
170     def report(self):
171         r = super().report()
172         bsz = max(self.metrics['bsz'], 1)
173         for k in ['know_loss', 'know_acc', 'know_chance']:
174             # round and average across all items since last report
175             r[k] = round_sigfigs(self.metrics[k] / bsz, 4)
176         return r
177 
178     def _parse_knowledge(self, obs):
179         if 'knowledge_parsed' in obs:
180             # make a copy of the list to prevent the future padding step from
181             # being destructive
182             return list(obs['knowledge_parsed'])
183 
184         checked_sentence = '{} {} {}'.format(
185             obs['title'], TOKEN_KNOWLEDGE, obs['checked_sentence']
186         )
187         # grab all the nonempty knowledge
188         obs_know = [k.strip() for k in obs.get('knowledge', '').split('\n')]
189         obs_know = [k for k in obs_know if k]
190 
191         # we want the correct knowledge to always be in index 0
192         try:
193             i = obs_know.index(checked_sentence)
194         except ValueError:
195             # uh oh, couldn't find the sentence in the knowledge. This happens for
196             # one or two examples in the training set. We can just artificially
197             # put it back in
198             i = 0
199             obs_know[0] = checked_sentence
200         obs_know[0], obs_know[i] = obs_know[i], obs_know[0]
201 
202         obs['knowledge_parsed'] = obs_know
203         obs['checked_sentence_parsed'] = checked_sentence
204         return obs['knowledge_parsed']
205 
206     def batchify(self, obs_batch):
207         """
208         Wizard custom batchify, which passes along the knowledge/title.
209 
210         Following the docstring of TorchAgent.batchify, it calls super, then
211         uses an extended version of the torch_agent.Batch namedtuple.
212 
213         The purpose of extending the info is to keep track of some custom
214         metrics.
215         """
216         batch = super().batchify(obs_batch)
217         reordered_observations = [obs_batch[i] for i in batch.valid_indices]
218         is_training = not ('eval_labels' in reordered_observations[0])
219 
220         # first parse and compile all the knowledge together
221         all_knowledges = []  # list-of-lists knowledge items for each observation
222         knowledge_counts = []  # how much knowledge each observation gets
223         for obs in reordered_observations:
224             obs_know = self._parse_knowledge(obs)
225             # downsample if desired
226             if (is_training and self.max_knowledge and
227                     len(obs_know) > self.max_knowledge):
228                 # offset by one so that we don't choose 0
229                 keepers = 1 + np.random.choice(
230                     len(obs_know) - 1, self.max_knowledge, False
231                 )
232                 # correct answer is always the first one
233                 keepers[0] = 0
234                 obs_know = [obs_know[i] for i in keepers]
235             all_knowledges.append(obs_know)
236             knowledge_counts.append(len(obs_know))
237 
238         # now we want to actually pack this into a tensor, along with the mask
239         N = len(reordered_observations)
240         K = max(knowledge_counts)
241         # round out the array so everything is equally sized
242         for i in range(N):
243             all_knowledges[i] += [''] * (K - knowledge_counts[i])
244         flattened_knowledge = list(chain(*all_knowledges))
245 
246         knowledge_vec = [
247             self._vectorize_text(
248                 # the beginning of the sentence is more useful
249                 k, truncate=self.knowledge_truncate, add_end=True, truncate_left=False,
250             )
251             for k in flattened_knowledge
252         ]
253         knowledge_vec, _ = padded_tensor(
254             knowledge_vec, self.NULL_IDX, self.use_cuda, left_padded=True,
255         )
256         knowledge_vec[:, -1] = self.END_IDX
257         T = knowledge_vec.size(-1)
258         knowledge_vec = knowledge_vec.view(N, K, T)
259 
260         # knowledge mask is a N x K tensor saying which items we're allowed to
261         # attend over
262         bsz = len(reordered_observations)
263         ck_mask = th.zeros(bsz, K).byte()
264         for i, klen in enumerate(knowledge_counts):
265             ck_mask[i, :klen] = 1
266         # and the correct labels
267         cs_ids = th.LongTensor(bsz).zero_()
268 
269         if self.use_cuda:
270             knowledge_vec = knowledge_vec.cuda()
271             ck_mask = ck_mask.cuda()
272             cs_ids = cs_ids.cuda()
273 
274         batch['know_vec'] = knowledge_vec
275         batch['ck_mask'] = ck_mask
276         batch['cs_ids'] = cs_ids
277         batch['use_cs_ids'] = is_training
278         batch['knowledge'] = np.array(flattened_knowledge).reshape(N, K)
279         return batch
280 
281     @classmethod
282     def add_cmdline_args(cls, argparser):
283         super(EndToEndAgent, cls).add_cmdline_args(argparser)
284         group = argparser.add_argument_group("EndToEnd Agent")
285         group.add_argument(
286             '--knowledge-alpha', type=float, default=0.95,
287             help='Weight on the knowledge-attn loss'
288         )
289         group.add_argument(
290             '--knowledge-truncate', type=int, default=32,
291             help='Knowledge truncation field. Defaults to same as --truncate.'
292         )
293         group.add_argument(
294             '--max-knowledge', type=int,
295             help='Reduce the amount of negative knowledge at train time.'
296         )
297         argparser.add_argument(
298             '--knowledge-alpha', type=float, default=0.95,
299             help='Weight on the knowledge-attn loss'
300         )
301 
302     def _model_input(self, batch):
303         return (
304             batch.text_vec,
305             batch.know_vec,
306             batch.ck_mask,
307             batch.cs_ids,
308             batch.use_cs_ids,
309         )
310 
311     def build_model(self):
312         self.model = EndToEndModel(self.opt, self.dict)
313         if self.opt['embedding_type'] != 'random':
314             self._copy_embeddings(
315                 self.model.embeddings.weight, self.opt['embedding_type']
316             )
317         if self.use_cuda:
318             self.model = self.model.cuda()
319         return self.model
320
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/projects/wizard_of_wikipedia/generator/agents.py b/projects/wizard_of_wikipedia/generator/agents.py
--- a/projects/wizard_of_wikipedia/generator/agents.py
+++ b/projects/wizard_of_wikipedia/generator/agents.py
@@ -58,48 +58,38 @@
             checked_sentences.append(checked_sentence)
 
         batch['checked_sentence'] = checked_sentences
+
         return batch
 
 
 class TwoStageAgent(_GenericWizardAgent):
-    def __init__(self, opt, shared):
+    def __init__(self, opt, shared=None):
         super().__init__(opt, shared)
         if shared is not None:
             # make sure the dialogue token appears
             self.dict[TOKEN_DIALOG] = 9999999
 
-    def observe(self, obs):
+    def _set_text_vec(self, obs, history, truncate):
         if 'text' not in obs:
             return obs
 
-        # TODO: resolve this with #1421
-        # get the dialog stuff
-        reply = self.last_reply()
-        self.observation = self.get_dialog_history(obs, reply=reply)
-        # we need to store the old text so that we can restore it
-        oldtext = obs['text']
-
-        # now we want to force prepend the knowledge stuff
-        fields = []
-        if 'chosen_topic' in obs:
-            fields += [obs['title']]
-        if 'checked_sentence' in obs:
-            fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]
-        if obs['text'] != '':
-            fields += [TOKEN_DIALOG, obs['text']]
-        obs['text'] = ' '.join(fields)
-
-        # now vectorize with the extra knowledge. It'll all get stored in the
-        # text_vec operation, etc
-        self.vectorize(
-            obs,
-            text_truncate=self.text_truncate,
-            label_truncate=self.label_truncate
-        )
-
-        # finally we need to return the old text to the way it was
-        obs['text'] = oldtext
-        assert obs is self.observation
+        if 'text_vec' not in obs:
+            fields = []
+            dialogue_history = history.get_history_str()
+            if 'chosen_topic' in obs:
+                fields += [obs['title']]
+            if 'checked_sentence' in obs:
+                fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]
+            if dialogue_history:
+                fields += [TOKEN_DIALOG, dialogue_history]
+            obs['text'] = ' '.join(fields)
+            obs['text_vec'] = self.dict.txt2vec(obs['text'])
+
+        # check truncation
+        if 'text_vec' in obs:
+            obs['text_vec'] = th.LongTensor(
+                self._check_truncate(obs['text_vec'], truncate, True)
+            )
 
         return obs

{"golden_diff": "diff --git a/projects/wizard_of_wikipedia/generator/agents.py b/projects/wizard_of_wikipedia/generator/agents.py\n--- a/projects/wizard_of_wikipedia/generator/agents.py\n+++ b/projects/wizard_of_wikipedia/generator/agents.py\n@@ -58,48 +58,38 @@\n             checked_sentences.append(checked_sentence)\n \n         batch['checked_sentence'] = checked_sentences\n+\n         return batch\n \n \n class TwoStageAgent(_GenericWizardAgent):\n-    def __init__(self, opt, shared):\n+    def __init__(self, opt, shared=None):\n         super().__init__(opt, shared)\n         if shared is not None:\n             # make sure the dialogue token appears\n             self.dict[TOKEN_DIALOG] = 9999999\n \n-    def observe(self, obs):\n+    def _set_text_vec(self, obs, history, truncate):\n         if 'text' not in obs:\n             return obs\n \n-        # TODO: resolve this with #1421\n-        # get the dialog stuff\n-        reply = self.last_reply()\n-        self.observation = self.get_dialog_history(obs, reply=reply)\n-        # we need to store the old text so that we can restore it\n-        oldtext = obs['text']\n-\n-        # now we want to force prepend the knowledge stuff\n-        fields = []\n-        if 'chosen_topic' in obs:\n-            fields += [obs['title']]\n-        if 'checked_sentence' in obs:\n-            fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]\n-        if obs['text'] != '':\n-            fields += [TOKEN_DIALOG, obs['text']]\n-        obs['text'] = ' '.join(fields)\n-\n-        # now vectorize with the extra knowledge. It'll all get stored in the\n-        # text_vec operation, etc\n-        self.vectorize(\n-            obs,\n-            text_truncate=self.text_truncate,\n-            label_truncate=self.label_truncate\n-        )\n-\n-        # finally we need to return the old text to the way it was\n-        obs['text'] = oldtext\n-        assert obs is self.observation\n+        if 'text_vec' not in obs:\n+            fields = []\n+            dialogue_history = history.get_history_str()\n+            if 'chosen_topic' in obs:\n+                fields += [obs['title']]\n+            if 'checked_sentence' in obs:\n+                fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]\n+            if dialogue_history:\n+                fields += [TOKEN_DIALOG, dialogue_history]\n+            obs['text'] = ' '.join(fields)\n+            obs['text_vec'] = self.dict.txt2vec(obs['text'])\n+\n+        # check truncation\n+        if 'text_vec' in obs:\n+            obs['text_vec'] = th.LongTensor(\n+                self._check_truncate(obs['text_vec'], truncate, True)\n+            )\n \n         return obs\n", "issue": "[Wizard of Wikipedia] Where is get_dialog_history() defined?\nhttps://github.com/facebookresearch/ParlAI/blob/a43f2880719c5a048fdf3d0aa5d5b25eeb9a1a41/projects/wizard_of_wikipedia/generator/agents.py#L78\r\n\r\n@stephenroller I'm guessing this got missed in TorchAgent. Could you please fix it?\n", "before_files": [{"content": "#!/usr/bin/env python\n\n# Copyright (c) Facebook, Inc. and its affiliates.\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nAgents for handling the generation aspects of Wizard.\n\"\"\"\nfrom itertools import chain\nfrom functools import lru_cache\n\nimport torch as th\nimport numpy as np\n\nfrom parlai.core.utils import padded_tensor, round_sigfigs\n\nfrom parlai.agents.transformer.transformer import TransformerGeneratorAgent\n\nfrom .modules import EndToEndModel\nfrom parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE\n\nTOKEN_DIALOG = '__dialog__'\n\n\nDEFAULT_OPTS = {\n    \"learningrate\": 5e-4,\n    \"optimizer\": \"adam\",\n    \"lr_scheduler\": \"invsqrt\",\n    \"warmup_updates\": 5000,\n    \"clip_norm\": 0.1,\n    \"ffn_size\": 512,\n    \"embedding_size\": 256,\n    \"n_heads\": 2,\n    \"dropout\": 0.2,\n    \"n_layers\": 5,\n    \"betas\": \"0.9,0.98\",\n    \"truncate\": 128,\n    \"dict_textfields\": \"text,labels,chosen_topic,checked_sentence,knowledge,title\",\n}\n\n\nclass _GenericWizardAgent(TransformerGeneratorAgent):\n    @classmethod\n    def add_cmdline_args(cls, argparser):\n        argparser.set_defaults(**DEFAULT_OPTS)\n        super(_GenericWizardAgent, cls).add_cmdline_args(argparser)\n\n    def batchify(self, obs_batch):\n        batch = super().batchify(obs_batch)\n        reordered_observations = [obs_batch[i] for i in batch.valid_indices]\n\n        checked_sentences = []\n        for obs in reordered_observations:\n            checked_sentence = '{} {} {}'.format(\n                obs.get('title', ''), TOKEN_KNOWLEDGE, obs.get('checked_sentence', '')\n            )\n            checked_sentences.append(checked_sentence)\n\n        batch['checked_sentence'] = checked_sentences\n        return batch\n\n\nclass TwoStageAgent(_GenericWizardAgent):\n    def __init__(self, opt, shared):\n        super().__init__(opt, shared)\n        if shared is not None:\n            # make sure the dialogue token appears\n            self.dict[TOKEN_DIALOG] = 9999999\n\n    def observe(self, obs):\n        if 'text' not in obs:\n            return obs\n\n        # TODO: resolve this with #1421\n        # get the dialog stuff\n        reply = self.last_reply()\n        self.observation = self.get_dialog_history(obs, reply=reply)\n        # we need to store the old text so that we can restore it\n        oldtext = obs['text']\n\n        # now we want to force prepend the knowledge stuff\n        fields = []\n        if 'chosen_topic' in obs:\n            fields += [obs['title']]\n        if 'checked_sentence' in obs:\n            fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]\n        if obs['text'] != '':\n            fields += [TOKEN_DIALOG, obs['text']]\n        obs['text'] = ' '.join(fields)\n\n        # now vectorize with the extra knowledge. It'll all get stored in the\n        # text_vec operation, etc\n        self.vectorize(\n            obs,\n            text_truncate=self.text_truncate,\n            label_truncate=self.label_truncate\n        )\n\n        # finally we need to return the old text to the way it was\n        obs['text'] = oldtext\n        assert obs is self.observation\n\n        return obs\n\n\nclass EndToEndAgent(_GenericWizardAgent):\n    def __init__(self, opt, shared=None):\n        super().__init__(opt, shared)\n        self._vectorize_text = lru_cache(int(2 ** 20))(self._vectorize_text)\n\n        # knowledge truncate defaults to the same as --truncate\n        self.knowledge_truncate = opt.get('knowledge_truncate')\n        if not self.knowledge_truncate:\n            self.knowledge_truncate = opt['truncate']\n        self.max_knowledge = opt.get('max_knowledge')\n        self.knowledge_alpha = opt['knowledge_alpha']\n\n    def _dummy_batch(self, bsz, maxlen):\n        batch = super()._dummy_batch(bsz, maxlen)\n        batch['know_vec'] = th.zeros(bsz, 2, 2).long().cuda()\n        batch['ck_mask'] = th.ones(bsz, 2).byte().cuda()\n        batch['cs_ids'] = th.zeros(bsz).long().cuda()\n        batch['use_cs_ids'] = True\n        return batch\n\n    def compute_loss(self, batch, return_output=False):\n        # first compute our regular forced decoding loss\n        token_loss, model_output = super().compute_loss(batch, return_output=True)\n        notnull = batch.label_vec.ne(self.NULL_IDX)\n        num_tokens = notnull.long().sum().item()\n\n        encoder_states = model_output[2]\n        ctx_know_attn = encoder_states[2]\n\n        if self.knowledge_alpha == 0.0:\n            loss = token_loss\n        else:\n            _, know_pred = ctx_know_attn.max(1)\n            know_acc = (know_pred == batch.cs_ids).float().sum().item()\n            know_chance = batch.ck_mask.sum(1).float().reciprocal().sum().item()\n            self.metrics['know_chance'] += know_chance\n            self.metrics['bsz'] += batch.text_vec.size(0)\n            self.metrics['know_acc'] += know_acc\n            know_loss = th.nn.functional.cross_entropy(\n                ctx_know_attn,\n                batch.cs_ids,\n                reduction='mean',\n            )\n            self.metrics['know_loss'] += know_loss.item() * batch.text_vec.size(0)\n            # in the original paper the loss was scaled by num_tokens for both\n            # know_loss and token_loss\n            know_loss /= num_tokens\n            loss = (\n                (1 - self.knowledge_alpha) * token_loss +\n                self.knowledge_alpha * know_loss\n            )\n        if return_output:\n            return (loss, model_output)\n        else:\n            return loss\n\n    def reset_metrics(self):\n        super().reset_metrics()\n        self.metrics['bsz'] = 0.0\n        self.metrics['know_acc'] = 0.0\n        self.metrics['know_loss'] = 0.0\n        self.metrics['know_chance'] = 0.0\n\n    def report(self):\n        r = super().report()\n        bsz = max(self.metrics['bsz'], 1)\n        for k in ['know_loss', 'know_acc', 'know_chance']:\n            # round and average across all items since last report\n            r[k] = round_sigfigs(self.metrics[k] / bsz, 4)\n        return r\n\n    def _parse_knowledge(self, obs):\n        if 'knowledge_parsed' in obs:\n            # make a copy of the list to prevent the future padding step from\n            # being destructive\n            return list(obs['knowledge_parsed'])\n\n        checked_sentence = '{} {} {}'.format(\n            obs['title'], TOKEN_KNOWLEDGE, obs['checked_sentence']\n        )\n        # grab all the nonempty knowledge\n        obs_know = [k.strip() for k in obs.get('knowledge', '').split('\\n')]\n        obs_know = [k for k in obs_know if k]\n\n        # we want the correct knowledge to always be in index 0\n        try:\n            i = obs_know.index(checked_sentence)\n        except ValueError:\n            # uh oh, couldn't find the sentence in the knowledge. This happens for\n            # one or two examples in the training set. We can just artificially\n            # put it back in\n            i = 0\n            obs_know[0] = checked_sentence\n        obs_know[0], obs_know[i] = obs_know[i], obs_know[0]\n\n        obs['knowledge_parsed'] = obs_know\n        obs['checked_sentence_parsed'] = checked_sentence\n        return obs['knowledge_parsed']\n\n    def batchify(self, obs_batch):\n        \"\"\"\n        Wizard custom batchify, which passes along the knowledge/title.\n\n        Following the docstring of TorchAgent.batchify, it calls super, then\n        uses an extended version of the torch_agent.Batch namedtuple.\n\n        The purpose of extending the info is to keep track of some custom\n        metrics.\n        \"\"\"\n        batch = super().batchify(obs_batch)\n        reordered_observations = [obs_batch[i] for i in batch.valid_indices]\n        is_training = not ('eval_labels' in reordered_observations[0])\n\n        # first parse and compile all the knowledge together\n        all_knowledges = []  # list-of-lists knowledge items for each observation\n        knowledge_counts = []  # how much knowledge each observation gets\n        for obs in reordered_observations:\n            obs_know = self._parse_knowledge(obs)\n            # downsample if desired\n            if (is_training and self.max_knowledge and\n                    len(obs_know) > self.max_knowledge):\n                # offset by one so that we don't choose 0\n                keepers = 1 + np.random.choice(\n                    len(obs_know) - 1, self.max_knowledge, False\n                )\n                # correct answer is always the first one\n                keepers[0] = 0\n                obs_know = [obs_know[i] for i in keepers]\n            all_knowledges.append(obs_know)\n            knowledge_counts.append(len(obs_know))\n\n        # now we want to actually pack this into a tensor, along with the mask\n        N = len(reordered_observations)\n        K = max(knowledge_counts)\n        # round out the array so everything is equally sized\n        for i in range(N):\n            all_knowledges[i] += [''] * (K - knowledge_counts[i])\n        flattened_knowledge = list(chain(*all_knowledges))\n\n        knowledge_vec = [\n            self._vectorize_text(\n                # the beginning of the sentence is more useful\n                k, truncate=self.knowledge_truncate, add_end=True, truncate_left=False,\n            )\n            for k in flattened_knowledge\n        ]\n        knowledge_vec, _ = padded_tensor(\n            knowledge_vec, self.NULL_IDX, self.use_cuda, left_padded=True,\n        )\n        knowledge_vec[:, -1] = self.END_IDX\n        T = knowledge_vec.size(-1)\n        knowledge_vec = knowledge_vec.view(N, K, T)\n\n        # knowledge mask is a N x K tensor saying which items we're allowed to\n        # attend over\n        bsz = len(reordered_observations)\n        ck_mask = th.zeros(bsz, K).byte()\n        for i, klen in enumerate(knowledge_counts):\n            ck_mask[i, :klen] = 1\n        # and the correct labels\n        cs_ids = th.LongTensor(bsz).zero_()\n\n        if self.use_cuda:\n            knowledge_vec = knowledge_vec.cuda()\n            ck_mask = ck_mask.cuda()\n            cs_ids = cs_ids.cuda()\n\n        batch['know_vec'] = knowledge_vec\n        batch['ck_mask'] = ck_mask\n        batch['cs_ids'] = cs_ids\n        batch['use_cs_ids'] = is_training\n        batch['knowledge'] = np.array(flattened_knowledge).reshape(N, K)\n        return batch\n\n    @classmethod\n    def add_cmdline_args(cls, argparser):\n        super(EndToEndAgent, cls).add_cmdline_args(argparser)\n        group = argparser.add_argument_group(\"EndToEnd Agent\")\n        group.add_argument(\n            '--knowledge-alpha', type=float, default=0.95,\n            help='Weight on the knowledge-attn loss'\n        )\n        group.add_argument(\n            '--knowledge-truncate', type=int, default=32,\n            help='Knowledge truncation field. Defaults to same as --truncate.'\n        )\n        group.add_argument(\n            '--max-knowledge', type=int,\n            help='Reduce the amount of negative knowledge at train time.'\n        )\n        argparser.add_argument(\n            '--knowledge-alpha', type=float, default=0.95,\n            help='Weight on the knowledge-attn loss'\n        )\n\n    def _model_input(self, batch):\n        return (\n            batch.text_vec,\n            batch.know_vec,\n            batch.ck_mask,\n            batch.cs_ids,\n            batch.use_cs_ids,\n        )\n\n    def build_model(self):\n        self.model = EndToEndModel(self.opt, self.dict)\n        if self.opt['embedding_type'] != 'random':\n            self._copy_embeddings(\n                self.model.embeddings.weight, self.opt['embedding_type']\n            )\n        if self.use_cuda:\n            self.model = self.model.cuda()\n        return self.model\n", "path": "projects/wizard_of_wikipedia/generator/agents.py"}], "after_files": [{"content": "#!/usr/bin/env python\n\n# Copyright (c) Facebook, Inc. and its affiliates.\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nAgents for handling the generation aspects of Wizard.\n\"\"\"\nfrom itertools import chain\nfrom functools import lru_cache\n\nimport torch as th\nimport numpy as np\n\nfrom parlai.core.utils import padded_tensor, round_sigfigs\n\nfrom parlai.agents.transformer.transformer import TransformerGeneratorAgent\n\nfrom .modules import EndToEndModel\nfrom parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE\n\nTOKEN_DIALOG = '__dialog__'\n\n\nDEFAULT_OPTS = {\n    \"learningrate\": 5e-4,\n    \"optimizer\": \"adam\",\n    \"lr_scheduler\": \"invsqrt\",\n    \"warmup_updates\": 5000,\n    \"clip_norm\": 0.1,\n    \"ffn_size\": 512,\n    \"embedding_size\": 256,\n    \"n_heads\": 2,\n    \"dropout\": 0.2,\n    \"n_layers\": 5,\n    \"betas\": \"0.9,0.98\",\n    \"truncate\": 128,\n    \"dict_textfields\": \"text,labels,chosen_topic,checked_sentence,knowledge,title\",\n}\n\n\nclass _GenericWizardAgent(TransformerGeneratorAgent):\n    @classmethod\n    def add_cmdline_args(cls, argparser):\n        argparser.set_defaults(**DEFAULT_OPTS)\n        super(_GenericWizardAgent, cls).add_cmdline_args(argparser)\n\n    def batchify(self, obs_batch):\n        batch = super().batchify(obs_batch)\n        reordered_observations = [obs_batch[i] for i in batch.valid_indices]\n\n        checked_sentences = []\n        for obs in reordered_observations:\n            checked_sentence = '{} {} {}'.format(\n                obs.get('title', ''), TOKEN_KNOWLEDGE, obs.get('checked_sentence', '')\n            )\n            checked_sentences.append(checked_sentence)\n\n        batch['checked_sentence'] = checked_sentences\n\n        return batch\n\n\nclass TwoStageAgent(_GenericWizardAgent):\n    def __init__(self, opt, shared=None):\n        super().__init__(opt, shared)\n        if shared is not None:\n            # make sure the dialogue token appears\n            self.dict[TOKEN_DIALOG] = 9999999\n\n    def _set_text_vec(self, obs, history, truncate):\n        if 'text' not in obs:\n            return obs\n\n        if 'text_vec' not in obs:\n            fields = []\n            dialogue_history = history.get_history_str()\n            if 'chosen_topic' in obs:\n                fields += [obs['title']]\n            if 'checked_sentence' in obs:\n                fields += [TOKEN_KNOWLEDGE, obs['checked_sentence']]\n            if dialogue_history:\n                fields += [TOKEN_DIALOG, dialogue_history]\n            obs['text'] = ' '.join(fields)\n            obs['text_vec'] = self.dict.txt2vec(obs['text'])\n\n        # check truncation\n        if 'text_vec' in obs:\n            obs['text_vec'] = th.LongTensor(\n                self._check_truncate(obs['text_vec'], truncate, True)\n            )\n\n        return obs\n\n\nclass EndToEndAgent(_GenericWizardAgent):\n    def __init__(self, opt, shared=None):\n        super().__init__(opt, shared)\n        self._vectorize_text = lru_cache(int(2 ** 20))(self._vectorize_text)\n\n        # knowledge truncate defaults to the same as --truncate\n        self.knowledge_truncate = opt.get('knowledge_truncate')\n        if not self.knowledge_truncate:\n            self.knowledge_truncate = opt['truncate']\n        self.max_knowledge = opt.get('max_knowledge')\n        self.knowledge_alpha = opt['knowledge_alpha']\n\n    def _dummy_batch(self, bsz, maxlen):\n        batch = super()._dummy_batch(bsz, maxlen)\n        batch['know_vec'] = th.zeros(bsz, 2, 2).long().cuda()\n        batch['ck_mask'] = th.ones(bsz, 2).byte().cuda()\n        batch['cs_ids'] = th.zeros(bsz).long().cuda()\n        batch['use_cs_ids'] = True\n        return batch\n\n    def compute_loss(self, batch, return_output=False):\n        # first compute our regular forced decoding loss\n        token_loss, model_output = super().compute_loss(batch, return_output=True)\n        notnull = batch.label_vec.ne(self.NULL_IDX)\n        num_tokens = notnull.long().sum().item()\n\n        encoder_states = model_output[2]\n        ctx_know_attn = encoder_states[2]\n\n        if self.knowledge_alpha == 0.0:\n            loss = token_loss\n        else:\n            _, know_pred = ctx_know_attn.max(1)\n            know_acc = (know_pred == batch.cs_ids).float().sum().item()\n            know_chance = batch.ck_mask.sum(1).float().reciprocal().sum().item()\n            self.metrics['know_chance'] += know_chance\n            self.metrics['bsz'] += batch.text_vec.size(0)\n            self.metrics['know_acc'] += know_acc\n            know_loss = th.nn.functional.cross_entropy(\n                ctx_know_attn,\n                batch.cs_ids,\n                reduction='mean',\n            )\n            self.metrics['know_loss'] += know_loss.item() * batch.text_vec.size(0)\n            # in the original paper the loss was scaled by num_tokens for both\n            # know_loss and token_loss\n            know_loss /= num_tokens\n            loss = (\n                (1 - self.knowledge_alpha) * token_loss +\n                self.knowledge_alpha * know_loss\n            )\n        if return_output:\n            return (loss, model_output)\n        else:\n            return loss\n\n    def reset_metrics(self):\n        super().reset_metrics()\n        self.metrics['bsz'] = 0.0\n        self.metrics['know_acc'] = 0.0\n        self.metrics['know_loss'] = 0.0\n        self.metrics['know_chance'] = 0.0\n\n    def report(self):\n        r = super().report()\n        bsz = max(self.metrics['bsz'], 1)\n        for k in ['know_loss', 'know_acc', 'know_chance']:\n            # round and average across all items since last report\n            r[k] = round_sigfigs(self.metrics[k] / bsz, 4)\n        return r\n\n    def _parse_knowledge(self, obs):\n        if 'knowledge_parsed' in obs:\n            # make a copy of the list to prevent the future padding step from\n            # being destructive\n            return list(obs['knowledge_parsed'])\n\n        checked_sentence = '{} {} {}'.format(\n            obs['title'], TOKEN_KNOWLEDGE, obs['checked_sentence']\n        )\n        # grab all the nonempty knowledge\n        obs_know = [k.strip() for k in obs.get('knowledge', '').split('\\n')]\n        obs_know = [k for k in obs_know if k]\n\n        # we want the correct knowledge to always be in index 0\n        try:\n            i = obs_know.index(checked_sentence)\n        except ValueError:\n            # uh oh, couldn't find the sentence in the knowledge. This happens for\n            # one or two examples in the training set. We can just artificially\n            # put it back in\n            i = 0\n            obs_know[0] = checked_sentence\n        obs_know[0], obs_know[i] = obs_know[i], obs_know[0]\n\n        obs['knowledge_parsed'] = obs_know\n        obs['checked_sentence_parsed'] = checked_sentence\n        return obs['knowledge_parsed']\n\n    def batchify(self, obs_batch):\n        \"\"\"\n        Wizard custom batchify, which passes along the knowledge/title.\n\n        Following the docstring of TorchAgent.batchify, it calls super, then\n        uses an extended version of the torch_agent.Batch namedtuple.\n\n        The purpose of extending the info is to keep track of some custom\n        metrics.\n        \"\"\"\n        batch = super().batchify(obs_batch)\n        reordered_observations = [obs_batch[i] for i in batch.valid_indices]\n        is_training = not ('eval_labels' in reordered_observations[0])\n\n        # first parse and compile all the knowledge together\n        all_knowledges = []  # list-of-lists knowledge items for each observation\n        knowledge_counts = []  # how much knowledge each observation gets\n        for obs in reordered_observations:\n            obs_know = self._parse_knowledge(obs)\n            # downsample if desired\n            if (is_training and self.max_knowledge and\n                    len(obs_know) > self.max_knowledge):\n                # offset by one so that we don't choose 0\n                keepers = 1 + np.random.choice(\n                    len(obs_know) - 1, self.max_knowledge, False\n                )\n                # correct answer is always the first one\n                keepers[0] = 0\n                obs_know = [obs_know[i] for i in keepers]\n            all_knowledges.append(obs_know)\n            knowledge_counts.append(len(obs_know))\n\n        # now we want to actually pack this into a tensor, along with the mask\n        N = len(reordered_observations)\n        K = max(knowledge_counts)\n        # round out the array so everything is equally sized\n        for i in range(N):\n            all_knowledges[i] += [''] * (K - knowledge_counts[i])\n        flattened_knowledge = list(chain(*all_knowledges))\n\n        knowledge_vec = [\n            self._vectorize_text(\n                # the beginning of the sentence is more useful\n                k, truncate=self.knowledge_truncate, add_end=True, truncate_left=False,\n            )\n            for k in flattened_knowledge\n        ]\n        knowledge_vec, _ = padded_tensor(\n            knowledge_vec, self.NULL_IDX, self.use_cuda, left_padded=True,\n        )\n        knowledge_vec[:, -1] = self.END_IDX\n        T = knowledge_vec.size(-1)\n        knowledge_vec = knowledge_vec.view(N, K, T)\n\n        # knowledge mask is a N x K tensor saying which items we're allowed to\n        # attend over\n        bsz = len(reordered_observations)\n        ck_mask = th.zeros(bsz, K).byte()\n        for i, klen in enumerate(knowledge_counts):\n            ck_mask[i, :klen] = 1\n        # and the correct labels\n        cs_ids = th.LongTensor(bsz).zero_()\n\n        if self.use_cuda:\n            knowledge_vec = knowledge_vec.cuda()\n            ck_mask = ck_mask.cuda()\n            cs_ids = cs_ids.cuda()\n\n        batch['know_vec'] = knowledge_vec\n        batch['ck_mask'] = ck_mask\n        batch['cs_ids'] = cs_ids\n        batch['use_cs_ids'] = is_training\n        batch['knowledge'] = np.array(flattened_knowledge).reshape(N, K)\n        return batch\n\n    @classmethod\n    def add_cmdline_args(cls, argparser):\n        super(EndToEndAgent, cls).add_cmdline_args(argparser)\n        group = argparser.add_argument_group(\"EndToEnd Agent\")\n        group.add_argument(\n            '--knowledge-alpha', type=float, default=0.95,\n            help='Weight on the knowledge-attn loss'\n        )\n        group.add_argument(\n            '--knowledge-truncate', type=int, default=32,\n            help='Knowledge truncation field. Defaults to same as --truncate.'\n        )\n        group.add_argument(\n            '--max-knowledge', type=int,\n            help='Reduce the amount of negative knowledge at train time.'\n        )\n        argparser.add_argument(\n            '--knowledge-alpha', type=float, default=0.95,\n            help='Weight on the knowledge-attn loss'\n        )\n\n    def _model_input(self, batch):\n        return (\n            batch.text_vec,\n            batch.know_vec,\n            batch.ck_mask,\n            batch.cs_ids,\n            batch.use_cs_ids,\n        )\n\n    def build_model(self):\n        self.model = EndToEndModel(self.opt, self.dict)\n        if self.opt['embedding_type'] != 'random':\n            self._copy_embeddings(\n                self.model.embeddings.weight, self.opt['embedding_type']\n            )\n        if self.use_cuda:\n            self.model = self.model.cuda()\n        return self.model\n", "path": "projects/wizard_of_wikipedia/generator/agents.py"}]}