text_plain_output_1.png
from flash.text import TextClassificationData, TextClassifier

import flash

import torch

datamodule = TextClassificationData.from_csv('review', 'sentiment', train_file='data/imdb/train.csv', val_file='data/imdb/valid.csv', batch_size=4)

model = TextClassifier(backbone='gchhablani/bert-base-cased-finetuned-sst2', labels=datamodule.labels)

trainer = flash.Trainer(max_epochs=2, gpus=torch.cuda.device_count())
trainer.finetune(model, datamodule=datamodule, strategy='freeze')