44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
|
import time
|
||
|
|
|
||
|
|
class DM:
|
||
|
|
def __init__(self):
|
||
|
|
self.utt_history = ""
|
||
|
|
self.history = []
|
||
|
|
self.state = None
|
||
|
|
|
||
|
|
def get_utt_history(self):
|
||
|
|
return self.utt_history
|
||
|
|
|
||
|
|
def get_recent_state(self):
|
||
|
|
return self.state
|
||
|
|
|
||
|
|
def get_dialogue_state(self):
|
||
|
|
return self.history
|
||
|
|
|
||
|
|
def update(self, new_state):
|
||
|
|
self.history.append(new_state)
|
||
|
|
self.utt_history = self.utt_history + " ||| " + new_state['modified_prompt']
|
||
|
|
self.state = {'intent': new_state['intent'],
|
||
|
|
'entities': new_state['entities'],
|
||
|
|
'offensive': new_state['is_offensive'],
|
||
|
|
'clear': new_state['is_clear'],
|
||
|
|
'time': time.time()}
|
||
|
|
|
||
|
|
def next_action(self):
|
||
|
|
if self.state['clear']:
|
||
|
|
if self.state['offensive']:
|
||
|
|
return "NoCanDo"
|
||
|
|
else:
|
||
|
|
if self.state['intent'] == 0:
|
||
|
|
return "RetGen"
|
||
|
|
elif self.state['intent'] == 1:
|
||
|
|
return "ConvGen"
|
||
|
|
elif self.state['intent'] == 2:
|
||
|
|
return "findPaper"
|
||
|
|
elif self.state['intent'] == 3:
|
||
|
|
return "findDataset"
|
||
|
|
elif self.state['intent'] == 4:
|
||
|
|
return "sumPaper"
|
||
|
|
else:
|
||
|
|
return "Clarify"
|