aitech-sd-lab/dialogue_system.py

from modules.NLU import NLU
from modules.DST import DST
from modules.DP import DP
from modules.NLG import NLG
import json
import re

value_dict = json.load(open('value_dict.json'))

def format_prediction(prediction, intent):
    out_list = []
    for idx, tup in enumerate(prediction):
        if tup[1][0] == 'B':
            slot_list = [intent, 'Cinema', tup[1][2:], tup[0]]
            for tup in prediction[idx + 1:]:
                if tup[1][0] != 'I':
                    break
                else:
                    slot_list[3] += ' ' + tup[0]
            out_list.append(slot_list)
    for slot in out_list:
        slot[3] = re.sub("^[!\"#$%&\'()*+,.;:<=>?\[\]^_`{|}~]+", '', slot[3])
        slot[3] = re.sub("[!\"#$%&\'()*+,.;:<=>?\[\]^_`{|}~]+$", '', slot[3])
    return out_list

def main():
    nlu = NLU()
    dst = DST()
    dp = DP()
    nlg = NLG()
    
    nlu.train_slot_model('data/train+test-pl.conllu', 'data/train+test-pl.conllu')
    nlu.train_intent_model('data/NLU_data_intent')
    # nlu.load_slot_model('slot-model-pl')
    # nlu.load_intent_model('intent-model-pl')

    print('===========================================')
    print('### By otrzymać pomoc, wpisz /pomoc ###')
    print('### By zakończyć rozmowę, wpisz /koniec ###')
    print('Witaj, jestem Usher - system do rezerwacji biletów kinowych. W czym mogę Ci pomóc?')
        
    # WIP
    while True:
        user_input = input('> ')
        user_input_lr = user_input.lower()
        flag = False
        slots = nlu.predict_slots(user_input)
        intent = nlu.predict_intent(user_input)
        formatted_prediction = format_prediction(slots, intent)
        for slot in formatted_prediction:
            if slot[2]=='seat':
                if ',' in slot[3]:  # ???
                    seat,row = slot[3].split(',')
                    formatted_prediction.remove(['inform', 'Cinema', 'seat', slot[3]])
                    formatted_prediction.append(['inform', 'Cinema', 'seat', seat])
                    formatted_prediction.append(['inform', 'Cinema', 'row', f' {row}'])
        
        if format_prediction == []:
            print('Czy mogę w czymś jeszcze pomóc? ')

        if user_input_lr == '/pomoc':
            print('TEKST_POMOCY_WIP')

        elif user_input_lr == '/koniec':
            print('Dziękuję za skorzystanie z moich usług. Miłego dnia!')
            break

        elif 'rezerw' in user_input_lr:
            if 'anulo' in user_input_lr:
                for slot in slots:
                    if slot == 'B-e-mail' or '@' in slot[0]:
                        print(nlg.update([['cinema','inform','cancel_book',slot[0]]]))
                        dst.update([['inform', 'Cinema', 'task', 'cancel_book'],['inform', 'Cinema', 'cancel_book_status',True]])
                        flag = True
                if not flag:
                    dst.update([['inform', 'Cinema', 'task', 'cancel_book'],['inform', 'Cinema', 'cancel_book_status',False]])
                    print(nlg.update([['cinema','inform','cancel_book','']]))

            else:
                dst.update([['inform', 'Cinema', 'task', 'book']])
                dst.update(formatted_prediction)
                for slot,value in dst.state['belief_state']['cinema']['book'].items():
                    if value == '':
                        print(nlg.update([['cinema','request',slot,'']]))
                        break
                    elif slot=="row" and value != '':
                        for slot,value in dst.state['belief_state']['cinema']['semi'].items():
                            if value == '':
                                print(nlg.update([['cinema','request',slot,'']]))
                                break   
                            elif slot in ["cancel_book_status","cancel_buy_status"] and value =='':
                                continue
        elif 'kup' in user_input_lr or ('zwr' and 'bilet') in user_input_lr:

            if 'anulo' in user_input_lr or 'zwr' in user_input_lr:
                for slot in slots:
                    if slot == 'B-e-mail' or '@' in slot[0]:
                        print(nlg.update([['cinema','inform','cancel_buy',slot[0]]]))
                        dst.update([['inform', 'Cinema', 'task', 'cancel_buy'],['inform', 'Cinema', 'cancel_buy_status',True]])
                        flag = True
                if not flag:
                    dst.update([['inform', 'Cinema', 'task', 'cancel_buy'],['inform', 'Cinema', 'cancel_buy_status',False]])
                    print(nlg.update([['cinema','inform','cancel_buy','']]))

            else:
                dst.update([['inform', 'Cinema', 'task', 'buy']])
                dst.update(formatted_prediction)
                print(dst.state)
                for slot,value in dst.state['belief_state']['cinema']['book'].items():
                    if value == '':
                        print(nlg.update([['cinema','request',slot,'']]))
                        break
                    elif slot=="row" and value != '':
                        for slot,value in dst.state['belief_state']['cinema']['semi'].items():
                            if value == '':
                                print(nlg.update([['cinema','request',slot,'']]))
                                break   
                            elif slot in ["cancel_book_status","cancel_buy_status"] and value =='':
                                continue
                      
        elif (('jak' or 'któr') and 'film') in user_input_lr or 'repertuar' in user_input_lr:   
            dst.update([['inform', 'Cinema', 'task', 'show_movies']])
            for slot in slots:
                if slot[0] in value_dict['train']['date'] or slot[0] in value_dict['train']['day'] or slot[0] in ['dziś','jutro','pojutrze']:
                    print(nlg.update([['cinema','inform','closestscreening',slot[0]]]))
                    flag = True
            if not flag:
                print(nlg.update([['cinema','inform','closestscreening','']]))

        elif (('czy' or 'jakie' or 'które') and ('dostęp' or 'woln' or 'zajęt') and 'miejsc') in user_input_lr:
            dst.update([['offer', 'Cinema', 'task', 'show_seats']])
            print(nlg.update([['cinema','offer','seat','']]))

        else:
            print(formatted_prediction) # NLU output
            dst.update(formatted_prediction)
            print(dst.state) # DST output 
            #DP
            for slot,value in dst.state['belief_state']['cinema']['book'].items():
                if value == '':
                    print(nlg.update([['cinema','request',slot,'']]))
                    break
                elif slot=="row" and value != '':
                    for slot,value in dst.state['belief_state']['cinema']['semi'].items():
                        if slot == 'payments' and value == '':
                            print('Czy chciałbyś dokonać płatności online?')
                            input_user = input('> ')
                            if input_user.lower() == 'tak':
                                print(nlg.update([['cinema','request',slot,'online']]))
                                dst.update([['inform', 'Cinema', 'payments', 'online']])
                            else:
                                print(nlg.update([['cinema','request',slot,'']]))
                                dst.update([['inform', 'Cinema', 'payments', 'in_cinema']])
                        elif value == '':
                            print(nlg.update([['cinema','request',slot,'']]))
                            break   
                        elif slot in ["cancel_book_status","cancel_buy_status"] and value =='':
                            continue
                       
        
            
    

if __name__ == '__main__':
    main()