ajsbsd commited on
Commit
4b6e90c
·
verified ·
1 Parent(s): e637501

Create run.py

Browse files
Files changed (1) hide show
  1. run.py +88 -0
run.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ !pip install -q transformers[torch] tokenizers datasets evaluate rouge_score sentencepiece huggingface_hub --upgrade
2
+
3
+ from huggingface_hub import notebook_login
4
+ notebook_login()
5
+
6
+ import nltk
7
+ from datasets import load_dataset
8
+ import evaluate
9
+ import numpy as np
10
+ from transformers import T5Tokenizer, DataCollatorForSeq2Seq
11
+ from transformers import T5ForConditionalGeneration, Seq2SeqTrainingArguments, Seq2SeqTrainer
12
+
13
+ # Load and split the dataset
14
+ dataset = load_dataset("ajsbsd/openbsd-faq")
15
+ dataset = dataset["train"].train_test_split(test_size=0.2)
16
+ #dataset = load_dataset("csv", data_files="./JEOPARDY_CSV.csv")
17
+ #dataset = dataset["train"].train_test_split(test_size=0.2)
18
+ # Load the tokenizer, model, and data collator
19
+ tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
20
+ model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
21
+ data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
22
+
23
+ # We prefix our tasks with "answer the question"
24
+ prefix = "Please answer this question: "
25
+
26
+ # Define our preprocessing function
27
+ def preprocess_function(examples):
28
+ """Add prefix to the sentences, tokenize the text, and set the labels"""
29
+ # The "inputs" are the tokenized answer:
30
+ inputs = [prefix + doc for doc in examples["question"]]
31
+ model_inputs = tokenizer(inputs, max_length=128, truncation=True)
32
+
33
+ # The "labels" are the tokenized outputs:
34
+ labels = tokenizer(text_target=examples["answer"], max_length=512, truncation=True)
35
+ model_inputs["labels"] = labels["input_ids"]
36
+ return model_inputs
37
+
38
+ # Map the preprocessing function across our dataset
39
+ tokenized_dataset = dataset.map(preprocess_function, batched=True)
40
+
41
+ # Set up Rouge score for evaluation
42
+ nltk.download("punkt", quiet=True)
43
+ metric = evaluate.load("rouge")
44
+
45
+ def compute_metrics(eval_preds):
46
+ preds, labels = eval_preds
47
+
48
+ # decode preds and labels
49
+ labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
50
+ decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
51
+ decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
52
+
53
+ # rougeLSum expects newline after each sentence
54
+ decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
55
+ decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
56
+
57
+ result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
58
+ return result
59
+
60
+ # Set up training arguments
61
+ training_args = Seq2SeqTrainingArguments(
62
+ output_dir="./flan-t5-base-openbsd-faq",
63
+ evaluation_strategy="epoch",
64
+ learning_rate=3e-4,
65
+ per_device_train_batch_size=8,
66
+ per_device_eval_batch_size=4,
67
+ weight_decay=0.01,
68
+ save_total_limit=3,
69
+ num_train_epochs=5,
70
+ predict_with_generate=True,
71
+ push_to_hub=False
72
+ )
73
+
74
+ # Set up trainer
75
+ trainer = Seq2SeqTrainer(
76
+ model=model,
77
+ args=training_args,
78
+ train_dataset=tokenized_dataset["train"],
79
+ eval_dataset=tokenized_dataset["test"],
80
+ tokenizer=tokenizer,
81
+ data_collator=data_collator,
82
+ compute_metrics=compute_metrics
83
+ )
84
+
85
+ # Train the model
86
+ trainer.train()
87
+
88
+ trainer.push_to_hub()