def test_serve():
model = TranslationTask(TEST_BACKBONE)
# TODO: Currently only servable once a preprocess and postprocess have been attached
model._preprocess = TranslationPreprocess(backbone=TEST_BACKBONE)
model._postprocess = Seq2SeqPostprocess()
model.eval()
model.serve()
def test_init_train(tmpdir):
if os.name == "nt":
# TODO: huggingface stuff timing out on windows
return True
model = TranslationTask(TEST_BACKBONE)
train_dl = torch.utils.data.DataLoader(DummyDataset())
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.fit(model, train_dl)
def test_jit(tmpdir):
sample_input = {
"input_ids": torch.randint(128, size=(1, 4)),
"attention_mask": torch.randint(1, size=(1, 4)),
}
path = os.path.join(tmpdir, "test.pt")
model = TranslationTask(TEST_BACKBONE, val_target_max_length=None)
model.eval()
# Huggingface only supports `torch.jit.trace`
model = torch.jit.trace(model, [sample_input])
torch.jit.save(model, path)
model = torch.jit.load(path)
out = model(sample_input)
assert isinstance(out, torch.Tensor)
from flash import download_data
from flash.text import TranslationData, TranslationTask
# 1. Download the data
download_data("https://pl-flash-data.s3.amazonaws.com/wmt_en_ro.zip", "data/")
# 2. Load the data
datamodule = TranslationData.from_files(train_file="data/wmt_en_ro/train.csv",
val_file="data/wmt_en_ro/valid.csv",
test_file="data/wmt_en_ro/test.csv",
input="input",
target="target",
batch_size=1)
# 3. Build the model
model = TranslationTask()
# 4. Create the trainer
trainer = flash.Trainer(precision=32,
gpus=int(torch.cuda.is_available()),
fast_dev_run=True)
# 5. Fine-tune the model
trainer.finetune(model, datamodule=datamodule)
# 6. Test model
trainer.test(model)
# 7. Save it!
trainer.save_checkpoint("translation_model_en_ro.pt")
backbone = "Helsinki-NLP/opus-mt-en-ro"
# 2. Load the data
datamodule = TranslationData.from_csv(
"input",
"target",
train_file="data/wmt_en_ro/train.csv",
val_file="data/wmt_en_ro/valid.csv",
test_file="data/wmt_en_ro/test.csv",
batch_size=1,
backbone=backbone,
)
# 3. Build the model
model = TranslationTask(backbone=backbone)
# 4. Create the trainer
trainer = flash.Trainer(
precision=16 if torch.cuda.is_available() else 32,
gpus=int(torch.cuda.is_available()),
fast_dev_run=True,
)
# 5. Fine-tune the model
trainer.finetune(model, datamodule=datamodule)
# 6. Test model
trainer.test(model)
# 7. Save it!
def test_init_train(tmpdir):
model = TranslationTask(TEST_BACKBONE)
train_dl = torch.utils.data.DataLoader(DummyDataset())
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
trainer.fit(model, train_dl)
from flash.core.data.utils import download_data
from flash.text import TranslationData, TranslationTask
# 1. Create the DataModule
download_data("https://pl-flash-data.s3.amazonaws.com/wmt_en_ro.zip", "./data")
datamodule = TranslationData.from_csv(
"input",
"target",
train_file="data/wmt_en_ro/train.csv",
val_file="data/wmt_en_ro/valid.csv",
backbone="Helsinki-NLP/opus-mt-en-ro",
)
# 2. Build the task
model = TranslationTask(backbone="Helsinki-NLP/opus-mt-en-ro")
# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())
trainer.finetune(model, datamodule=datamodule)
# 4. Translate something!
predictions = model.predict([
"BBC News went to meet one of the project's first graduates.",
"A recession has come as quickly as 11 months after the first rate hike and as long as 86 months.",
"Of course, it's still early in the election cycle.",
])
print(predictions)
# 5. Save the model!
trainer.save_checkpoint("translation_model_en_ro.pt")
def test_serve():
model = TranslationTask(TEST_BACKBONE)
model.eval()
model.serve()
# 1. Download the data
if args.download:
download_data("https://pl-flash-data.s3.amazonaws.com/wmt_en_ro.zip",
"data/")
# 2. Load the data
datamodule = TranslationData.from_files(
train_file="data/wmt_en_ro/train.csv",
valid_file="data/wmt_en_ro/valid.csv",
test_file="data/wmt_en_ro/test.csv",
input="input",
target="target",
)
# 3. Build the model
model = TranslationTask(backbone=args.backbone)
# 4. Create the trainer. Run once on data
trainer = flash.Trainer(max_epochs=args.max_epochs,
precision=16,
gpus=args.gpus)
# 5. Fine-tune the model
trainer.finetune(model, datamodule=datamodule)
# 6. Test model
trainer.test()
# 7. Save it!
trainer.save_checkpoint("translation_model_en_ro.pt")