Views
No views yet
1from transformers import AutoTokenizer, AutoModelForQuestionAnswering
2
3tokenizer = AutoTokenizer.from_pretrained("Primer/bart-squad2")
4model = AutoModelForQuestionAnswering.from_pretrained("Primer/bart-squad2")
5model.to('cuda'); model.eval()
6
7def answer(question, text):
8 seq = '<s>' + question + ' </s> </s> ' + text + ' </s>'
9 tokens = tokenizer.encode_plus(seq, return_tensors='pt', padding='max_length', max_length=1024)
10 input_ids = tokens['input_ids'].to('cuda')
11 attention_mask = tokens['attention_mask'].to('cuda')
12 start, end, _ = model(input_ids, attention_mask=attention_mask)
13 start_idx = int(start.argmax().int())
14 end_idx = int(end.argmax().int())
15 print(tokenizer.decode(input_ids[0, start_idx:end_idx]).strip())
16 # ^^ it will be an empty string if the model decided "unanswerable"
17
18>>> question = "Where does Tom live?"
19>>> context = "Tom is an engineer in San Francisco."
20>>> answer(question, context)
21San Francisco.to('cuda') stuff if running on CPU).run_squad.py with:| param | value |
|---|---|
| batch size | 8 |
| max_seq_length | 1024 |
| learning rate | 1e-5 |
| epochs | 2 |