INNER CODE UNIT · Python
process_goal
princeton-nlp/WebShop · transfer/app.py:34
def process_goal(state):
state = state.lower().replace('"', '').replace("'", "")
state = state.replace('amazon shopping game\ninstruction:', '').replace('webshop\ninstruction:', '')
state = state.replace('\n[button] search [button_]', '').strip()
if ', and price lower than' in state:
state = state.split(', and price lower than')[0]
return state
def data_collator(batch):
state_input_ids, state_attention_mask, action_input_ids, action_attention_mask, sizes, labels, images = [], [], [], [], [], [], []
for sample in batch:
state_input_ids.append(sample['state_input_ids'])
state_attention_mask.append(sample['state_attention_mask'])
action_input_ids.extend(sample['action_input_ids'])
action_attention_mask.extend(sample['action_attention_mask'])
sizes.append(sample['sizes'])
labels.append(sample['labels'])