Views
No views yet
{base_model}-{target_dataset}teabreac-{base_model}teabreac-{base_model}-{target_dataset}base_model above can be from: bart-large, t5-large, t5-3b, nt5-small, preasm-large.
The target_dataset above can be from: drop, tatqa, iirc-gold, iirc-retrieved, numglue.1# NOTE: This model is only pretrained on TeaBReaC, and not on any real QA dataset.
2from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3from digit_tokenization import enable_digit_tokenization # digit_tokenization.py from https://github.com/stonybrooknlp/teabreac
4
5model_name = "StonyBrookNLP/teabreac-bart-large"
6tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False) # Fast doesn't work with digit tokenization
7model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
8enable_digit_tokenization(tokenizer)
9input_texts = [
10 "answer_me: Who scored the first touchdown of the game?" +
11 "context: ... Oakland would get the early lead in the first quarter as quarterback JaMarcus Russell completed a 20-yard touchdown pass to rookie wide receiver Chaz Schilens..."
12 # Note: some models have slightly different qn/ctxt format. See the github repo.
13]
14input_ids = tokenizer(
15 input_texts, return_tensors="pt",
16 truncation=True, max_length=800,
17 add_special_tokens=True, padding=True,
18)["input_ids"]
19generated_ids = model.generate(input_ids, min_length=1, max_length=50)
20generated_predictions = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)
21generated_predictions = [
22 tokenizer.fix_decoded_text(generated_prediction) for generated_prediction in generated_predictions
23]
24# => ["Chaz Schilens"]