def test_knowledge_retriever(self):
from parlai.core.params import ParlaiParser
parser = ParlaiParser(False, False)
KnowledgeRetrieverAgent.add_cmdline_args(parser)
parser.set_params(
model='projects:wizard_of_wikipedia:knowledge_retriever',
add_token_knowledge=True,
)
knowledge_opt = parser.parse_args([], print_args=False)
knowledge_agent = create_agent(knowledge_opt)
knowledge_agent.observe({
'text': 'what do you think of mountain dew?',
'chosen_topic': 'Mountain Dew',
'episode_done': False,
})
knowledge_act = knowledge_agent.act()
title = knowledge_act['title']
self.assertEqual(title, 'Mountain Dew',
'Did not save chosen topic correctly')
knowledge = knowledge_act['text']
self.assertIn(TOKEN_KNOWLEDGE, knowledge,
'Knowledge token was not inserted correctly')
checked_sentence = knowledge_act['checked_sentence']
self.assertEqual(
checked_sentence,
'Mountain Dew (stylized as Mtn Dew) is a carbonated soft drink brand produced and owned by PepsiCo.',
'Did not correctly choose the checked sentence',
)
def create_retriever():
parser = ParlaiParser(False, False)
KnowledgeRetrieverAgent.add_cmdline_args(parser)
parser.set_params(
model='projects:wizard_of_wikipedia:knowledge_retriever',
add_token_knowledge=False,
)
knowledge_opt = parser.parse_args([])
return create_agent(knowledge_opt)
def _set_up_knowledge_agent(self, add_token_knowledge=False):
from parlai.core.params import ParlaiParser
parser = ParlaiParser(False, False)
KnowledgeRetrieverAgent.add_cmdline_args(parser, partial_opt=self.opt)
parser.set_params(
model='projects:wizard_of_wikipedia:knowledge_retriever',
add_token_knowledge=add_token_knowledge,
)
knowledge_opt = parser.parse_args([])
self.knowledge_agent = create_agent(knowledge_opt)
def _set_up_knowledge_agent(self,
add_token_knowledge: bool = False,
shared=None) -> None:
"""
Set up knowledge agent for knowledge retrieval generated from WoW project.
"""
parser = ParlaiParser(False, False)
KnowledgeRetrieverAgent.add_cmdline_args(parser)
parser.set_params(
model='projects:wizard_of_wikipedia:knowledge_retriever',
add_token_knowledge=add_token_knowledge,
)
knowledge_opt = parser.parse_args([])
if shared:
self.knowledge_agent = KnowledgeRetrieverAgent(
knowledge_opt, shared.get('knowledge_retriever', None))
else:
self.knowledge_agent = KnowledgeRetrieverAgent(knowledge_opt)