fix sharing
NOTE: this model only works for eval, it assumes all training is already done.
"""
-from parlai.core.agents import Agent, create_agent
+from parlai.core.agents import Agent, create_agent, create_agent_from_shared
from projects.wizard_of_wikipedia.wizard_transformer_ranker.wizard_transformer_ranker import (
WizardTransformerRankerAgent,
)
@@ -34,15 +34,15 @@ class InteractiveRetrievalAgent(Agent):
'full_dialogue_retrieval_model',
)
+ if not shared:
# Create responder
self._set_up_responder(opt)
-
- if not shared:
# Create retriever
self._set_up_retriever(opt)
else:
self.opt = shared['opt']
self.retriever = shared['retriever']
+ self.responder = create_agent_from_shared(shared['responder_shared_opt'])
self.sent_tok = shared['sent_tok']
self.wiki_map = shared['wiki_map']
@@ -288,6 +288,7 @@ class InteractiveRetrievalAgent(Agent):
shared = super().share()
shared['opt'] = self.opt
shared['retriever'] = self.retriever
+ shared['responder_shared_opt'] = self.responder.share()
shared['sent_tok'] = self.sent_tok
shared['wiki_map'] = self.wiki_map
