Quick Start

Training an agent on the knowledge-grounded dialogue task

Import the modules

from cogagent import *
import torch
import torch.nn as nn
import torch.optim as optim

Init the logger, device and experiment result saving dir

device, output_path = init_cogagent(
	device_id=8,
	output_path=datapath,
	folder_tag="run_diffks_on_wow",
)

Choose utterance reader

reader = WoWReader(raw_data_path=raw_data_path)
train_data, dev_data, test_data = reader.read_all()
vocab = reader.read_vocab()

Choose data processor

 
# In the training phase, no retriever is selected as the knowledge is provided by dataset
processor = WoWForDiffksProcessor(max_token_len=512, vocab=vocab, debug=False)
train_dataset = processor.process_train(train_data)
dev_dataset = processor.process_dev(dev_data)
test_dataset = processor.process_test(test_data)

Choose response generator

 
model = DiffKSModel()
metric = BaseKGCMetric(default_metric_name="bleu-4",vocab=vocab)
loss = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)

Start Training

# Use the provided Trainer class to start the model training process
trainer = Trainer(model,train_dataset,dev_data=test_dataset,n_epochs=40,batch_size=2,
					loss=loss,optimizer=optimizer,scheduler=None,metrics=metric,
					drop_last=False,gradient_accumulation_steps=1,num_workers=5,
					validate_steps=2000,save_by_metric="bleu-4",save_steps=None,
					output_path=output_path,grad_norm=1,
					use_tqdm=True,device=device,
					fp16_opt_level='O1',
					)
trainer.train()