JARVIS/dialog_model.py

from typing import Any
import jsgf
import argparse


class Model():
    def __init__(self):
        pass

    def __call__(self, prompt) -> Any:
        nlu = NLU()
        dst = DST()
        dp = DP()
        nlg = NLG()

        msg = prompt.lower()

        rama_nlu = nlu(msg)
        print(rama_nlu)
        rama_dst = dst(rama_nlu)
        rama_dp = dp(rama_dst)
        text = nlg(rama_dp)

        return text


class NLU():
    def __init__(self):
        pass

    def get_dialog_act(self, rule):
        slots = []
        self.get_slots(rule.expansion, slots)
        return {'act': rule.grammar.name, 'slots': slots}

    def get_slots(self, expansion, slots):
        if expansion.tag != '':
            slots.append((expansion.tag, expansion.current_match))
            return

        for child in expansion.children:
            self.get_slots(child, slots)

        if not expansion.children and isinstance(expansion, jsgf.NamedRuleRef):
            self.get_slots(expansion.referenced_rule.expansion, slots)

    def __call__(self, prompt) -> Any:
        book_grammar = jsgf.parse_grammar_file('book.jsgf')

        matched = book_grammar.find_matching_rules(prompt)

        if matched:
            return self.get_dialog_act(matched[0])
        else:
            return {'act': 'null', 'slots': []}


class DST():
    def __init__(self):
        self.msgs = []

    def __call__(self, msg) -> Any:
        self.msgs.append(msg)
        return msg


class DP():
    def __init__(self):
        pass

    def __call__(self, msg) -> Any:
        if "imie" in msg:
            return "imieMSG"
        else:
            return None


class NLG():
    def __init__(self):
        pass

    def __call__(self, msg) -> Any:
        if msg == "imieMSG":
            return "Mam na imie JARVIS"
        else:
            return "Nie rozumiem"


if __name__ == "__main__":
    model = Model()
    parser = argparse.ArgumentParser()
    parser.add_argument("--msg")
    args = parser.parse_args()

    ans = model(prompt=args.msg)
    print(ans)