File size: 4,209 Bytes
52943d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import json
import random
from typing import List, Dict, Tuple
from adaptor.evaluators.generative import ROUGE, BLEU
from adaptor.lang_module import LangModule
from adaptor.objectives.seq2seq import Sequence2Sequence
from adaptor.utils import AdaptationArguments, StoppingStrategy
from adaptor.schedules import ParallelSchedule
from adaptor.adapter import Adapter
import wandb
# Dataset creation
## Define paths to JSON files
db_path = 'db_schemas.json'
spider_dataset_train_path = 'spider/train_spider.json'
spider_dataset_dev_path = 'spider/dev.json'
spider_syn_train_path = 'Spider-Syn/train_spider.json'
spider_syn_dev_path = 'Spider-Syn/dev.json'
## Open files
with open(db_path, 'r') as file_db:
database_schemas = json.load(file_db)
with open(spider_dataset_train_path, 'r') as file_spider:
spider_train_dataset = json.load(file_spider)
with open(spider_dataset_dev_path, 'r') as file_spider:
spider_dev_dataset = json.load(file_spider)
with open(spider_syn_train_path, 'r') as file_spider:
spider_syn_train_dataset = json.load(file_spider)
with open(spider_syn_dev_path, 'r') as file_spider:
spider_syn_dev_dataset = json.load(file_spider)
## Include spider questions with synonyms (questions include text which is not in DB columns)
spider_train_dataset.extend([question for question in spider_syn_train_dataset if question['SpiderQuestion']!=question['SpiderSynQuestion']])
spider_dev_dataset.extend([question for question in spider_syn_dev_dataset if question['SpiderQuestion']!=question['SpiderSynQuestion']])
random.shuffle(spider_train_dataset)
random.shuffle(spider_dev_dataset)
def create_prompt(question: str, schema: str) -> str:
return " ".join(["Question: ",question, "Schema:", schema])
def create_vals_and_labels(dataset: List[dict], db_dict: Dict[str, str]) -> Tuple[List[str], List[str]]:
list_labels = [data["query"] for data in dataset]
list_prompts = [create_prompt(data["question"], db_dict[data["db_id"]])
if "question" in data else create_prompt(data["SpiderSynQuestion"], db_dict[data["db_id"]]) for data in dataset]
return list_prompts, list_labels
## Training prompts and labels
prompts_train, labels_train = create_vals_and_labels(spider_train_dataset, database_schemas)
assert len(prompts_train)==len(labels_train)
## Evaluation prompts and labels
prompts_eval, labels_eval = create_vals_and_labels(spider_dev_dataset, database_schemas)
assert len(prompts_eval)==len(labels_eval)
# Training
lang_module = LangModule("google/t5-large-lm-adapt")
evaluators = [BLEU(), ROUGE(decides_convergence=True)]
wandb.init(project="chatbot")
seq_qa = Sequence2Sequence(lang_module,
texts_or_path=prompts_train,
labels_or_path=labels_train,
val_texts_or_path=prompts_eval,
val_labels_or_path=labels_eval,
batch_size=4,
val_evaluators=evaluators,
objective_id="txt2SQL_Spider")
training_arguments = AdaptationArguments(output_dir="checkpoints-txt2sql",
learning_rate=5e-5,
stopping_strategy=StoppingStrategy.ALL_OBJECTIVES_CONVERGED,
stopping_patience=8,
save_total_limit=8,
do_train=True,
do_eval=True,
bf16=True,
warmup_steps=100,
gradient_accumulation_steps=8,
logging_steps=10,
eval_steps=200,
save_steps=200,
num_train_epochs=10,
evaluation_strategy="steps")
schedule = ParallelSchedule(objectives=[seq_qa],
args=training_arguments)
adapter = Adapter(lang_module, schedule, args=training_arguments)
adapter.train() |