INNER CODE UNIT · Python
bart_predict
princeton-nlp/WebShop · transfer/app.py:66
def bart_predict(input):
input_ids = bart_tokenizer(input)['input_ids']
input_ids = torch.tensor(input_ids).unsqueeze(0)
output = bart_model.generate(input_ids, max_length=512, num_return_sequences=5, num_beams=5)
return bart_tokenizer.batch_decode(output.tolist(), skip_special_tokens=True)[0]
def bert_predict(obs, info, softmax=True):
valid_acts = info['valid']
assert valid_acts[0].startswith('click[')
state_encodings = bert_tokenizer(process_str(obs), max_length=512, truncation=True, padding='max_length')
action_encodings = bert_tokenizer(list(map(process_str, valid_acts)), max_length=512, truncation=True, padding='max_length')
batch = {
'state_input_ids': state_encodings['input_ids'],
'state_attention_mask': state_encodings['attention_mask'],
'action_input_ids': action_encodings['input_ids'],
'action_attention_mask': action_encodings['attention_mask'],
'sizes': len(valid_acts),