mel_agent.py

'''
mel_agent.py - An example dialogue agent class
==========================================================================

Build up an pipeline agent with nlu, dst, policy and nlg.

@author: Songbo
'''


from convlab2.dialog_agent.agent import DialogueAgent
from convlab2.nlu.jointBERT.multiwoz import BERTNLU
from convlab2.dst.rule.multiwoz import RuleDST
from convlab2.policy.mel import MEL
from convlab2.nlg.template.multiwoz import TemplateNLG


class Agent(DialogueAgent):

        def __init__(self):
            nlu = BERTNLU()
            dst = RuleDST()

            model_path = "convlab2/policy/best_policies/best_supervised_MEL_model.ckpt"
            policy = MEL()
            policy.load(model_path)

            nlg = TemplateNLG(is_user=False)
            super().__init__(nlu, dst, policy, nlg)

            self.agent_name = "mel_agent"

# END OF FILE