text_plain_output_4.png text_plain_output_3.png text_plain_output_2.png 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')

datamodule = TextClassificationData.from_lists(predict_data=["Joker's performance was outstanding!!!", 'This is the best movie ever!!!', 'This movie was terrible what a waste of time'], batch_size=10)

reloaded_model = TextClassifier.load_from_checkpoint('text_classification_model.pt')

flash.Trainer().predict(reloaded_model, datamodule=datamodule, output='labels')