review_sentiment_flow.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

import os

from metaflow import (
    FlowSpec,
    IncludeFile,
    Parameter,
    card,
    current,
    step,
    environment,
    kubernetes,
    pypi,
    nvidia,
    retry,
)

GCS_PROJECT_NAME = "moz-fx-mlops-inference-nonprod"
GCS_BUCKET_NAME = "mf-models-test1"
MODEL_STORAGE_PATH = "ctroy-example-flow/model-bytes.pth"
TOKENIZER_STORAGE_PATH = "models/tokenizer/"

class ReviewSentimentFlow(FlowSpec):
    """
    A sample flow demonstrating
    The use of custom docker images and GPU facilities
    to train a biggish machine learning model (i.e. not a toy example on the iris dataset).
    """

    # This is an example of a parameter. You can toggle this when you call the flow
    # with python template_flow.py run --offline False
    offline_wandb = Parameter(
        "offline",
        help="Do not connect to W&B servers when training",
        type=bool,
        default=True,
    )

    # You can import the contents of files from your file system to use in flows.
    # This is meant for small files—in this example, a bit of config.
    example_config = IncludeFile("example_config", default="./example_config.json")

    @card(type="default")
    @kubernetes
    @step
    def start(self):
        """
        Each flow has a 'start' step.

        You can use it for collecting/preprocessing data or other setup tasks.
        """

        self.next(self.train_model)

    @card
    @environment(
        vars={
            "WANDB_API_KEY": os.getenv("WANDB_API_KEY"),
            "WANDB_PROJECT": os.getenv("WANDB_PROJECT"),
        }
    )
    @pypi(python='3.10.8',
          packages={
              'torch': '2.4.1',
              'wandb': '0.17.8',
              'datasets': '3.0.0',
              'numpy': '1.26.4',
              'tqdm': '4.66.5',
              'transformers': '4.44.2',
              'mozmlops': '0.1.4',
          })
    @nvidia
    @retry
    @step
    def train_model(self):
        """
        Trains a transformer model on movie reviews
        using NVIDIA GPUs
        """
        import json
        import wandb
        import collections

        import datasets
        import numpy as np
        import torch
        import torch.nn as nn
        import torch.optim as optim
        import tqdm
        import transformers
        import json

        from io import BytesIO
        from review_sentiment_model import ReviewSentimentModel
        from os.path import isfile, join

        config_as_dict = json.loads(self.example_config)
        print(f"The config file says: {config_as_dict.get('example_key')}")

        if not self.offline_wandb:
            tracking_run = wandb.init(project=os.getenv("WANDB_PROJECT"))
            wandb_url = tracking_run.get_url()
            current.card.append(Markdown("# Weights & Biases"))
            current.card.append(
                Markdown(f"Your training run is tracked [here]({wandb_url}).")
            )

        print("All set. Running training.")

        seed = 1234

        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.backends.cudnn.deterministic = True

        train_data, test_data = datasets.load_dataset("imdb", split=["train", "test"])
        transformer_name = "bert-base-uncased"
        tokenizer = transformers.AutoTokenizer.from_pretrained(transformer_name)

        print(tokenizer.tokenize("hello world!"))
        print(tokenizer.encode("hello world!"))
        print(tokenizer.convert_ids_to_tokens(tokenizer.encode("hello world")))
        print(tokenizer("hello world!"))

        def tokenize_and_numericalize_example(example, tokenizer):
            ids = tokenizer(example["text"], truncation=True)["input_ids"]
            return {"ids": ids}

        train_data = train_data.map(
            tokenize_and_numericalize_example, fn_kwargs={"tokenizer": tokenizer}
        )
        test_data = test_data.map(
            tokenize_and_numericalize_example, fn_kwargs={"tokenizer": tokenizer}
        )

        print(train_data[0])

        test_size = 0.25
        pad_index = tokenizer.pad_token_id

        train_valid_data = train_data.train_test_split(test_size=test_size)
        train_data = train_valid_data["train"]
        valid_data = train_valid_data["test"]

        train_data = train_data.with_format(type="torch", columns=["ids", "label"])
        valid_data = valid_data.with_format(type="torch", columns=["ids", "label"])
        test_data = test_data.with_format(type="torch", columns=["ids", "label"])

        def get_collate_fn(pad_index):
            def collate_fn(batch):
                batch_ids = [i["ids"] for i in batch]
                batch_ids = nn.utils.rnn.pad_sequence(
                    batch_ids, padding_value=pad_index, batch_first=True
                )
                batch_label = [i["label"] for i in batch]
                batch_label = torch.stack(batch_label)
                batch = {"ids": batch_ids, "label": batch_label}
                return batch

            return collate_fn

        def get_data_loader(dataset, batch_size, pad_index, shuffle=False):
            collate_fn = get_collate_fn(pad_index)
            data_loader = torch.utils.data.DataLoader(
                dataset=dataset,
                batch_size=batch_size,
                collate_fn=collate_fn,
                shuffle=shuffle,
            )
            return data_loader

        batch_size = 8

        train_data_loader = get_data_loader(train_data, batch_size, pad_index, shuffle=True)
        valid_data_loader = get_data_loader(valid_data, batch_size, pad_index)
        test_data_loader = get_data_loader(test_data, batch_size, pad_index)

        transformer = transformers.AutoModel.from_pretrained(transformer_name)
        self.output_dim = len(train_data["label"].unique())
        freeze = False

        model = ReviewSentimentModel(transformer, self.output_dim, freeze)
        lr = 1e-5

        optimizer = optim.Adam(model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        model = model.to(device)
        criterion = criterion.to(device)

        def train(data_loader, model, criterion, optimizer, device):
            model.train()
            epoch_losses = []
            epoch_accs = []
            for batch in tqdm.tqdm(data_loader, desc="training..."):
                ids = batch["ids"].to(device)
                label = batch["label"].to(device)
                prediction = model(ids)
                loss = criterion(prediction, label)
                accuracy = get_accuracy(prediction, label)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                epoch_losses.append(loss.item())
                epoch_accs.append(accuracy.item())
            return np.mean(epoch_losses), np.mean(epoch_accs)

        def evaluate(data_loader, model, criterion, device):
            model.eval()
            epoch_losses = []
            epoch_accs = []
            with torch.no_grad():
                for batch in tqdm.tqdm(data_loader, desc="evaluating..."):
                    ids = batch["ids"].to(device)
                    label = batch["label"].to(device)
                    prediction = model(ids)
                    loss = criterion(prediction, label)
                    accuracy = get_accuracy(prediction, label)
                    epoch_losses.append(loss.item())
                    epoch_accs.append(accuracy.item())
            return np.mean(epoch_losses), np.mean(epoch_accs)

        def get_accuracy(prediction, label):
            batch_size, _ = prediction.shape
            predicted_classes = prediction.argmax(dim=-1)
            correct_predictions = predicted_classes.eq(label).sum()
            accuracy = correct_predictions / batch_size
            return accuracy

        n_epochs = 3
        best_valid_loss = float("inf")

        metrics = collections.defaultdict(list)

        for epoch in range(n_epochs):
            train_loss, train_acc = train(
                train_data_loader, model, criterion, optimizer, device
            )
            valid_loss, valid_acc = evaluate(valid_data_loader, model, criterion, device)
            metrics["train_losses"].append(train_loss)
            metrics["train_accs"].append(train_acc)
            metrics["valid_losses"].append(valid_loss)
            metrics["valid_accs"].append(valid_acc)
            if valid_loss < best_valid_loss:
                best_valid_loss = valid_loss
                torch.save(model.state_dict(), "transformer.pt")
            print(f"epoch: {epoch}")
            print(f"train_loss: {train_loss:.3f}, train_acc: {train_acc:.3f}")
            print(f"valid_loss: {valid_loss:.3f}, valid_acc: {valid_acc:.3f}")

        model.load_state_dict(torch.load("transformer.pt"))

        test_loss, test_acc = evaluate(test_data_loader, model, criterion, device)
        print(f"test_loss: {test_loss:.3f}, test_acc: {test_acc:.3f}")

        self.device = device
        buffer = BytesIO()
        torch.save(model.state_dict(), buffer)
        self.model_state_dict_bytes = buffer.getvalue()

        self.tokenizer_as_dict = {}
        tokenizer.save_pretrained("tok")

        with open('tok/tokenizer.json') as file:
            self.tokenizer_as_dict = json.load(file)

        self.next(self.error_analysis)

    @pypi(python='3.10.8',
          packages={
              'torch': '2.4.1',
              'wandb': '0.17.8',
              'tokenizers' : '0.20.0',
          })
    @kubernetes
    @step
    def error_analysis(self):
        """
        Predict the sentiment of some sample movie reviews and see,
        on an individual level, how they look
        """
        import torch
        from tokenizers import Tokenizer
        from review_sentiment_model import ReviewSentimentModel

        from io import BytesIO
        import json

        device = self.device

        with open('tokenizer.json', 'w') as fp:
            json.dump(self.tokenizer_as_dict, fp)

        tokenizer = Tokenizer.from_file("tokenizer.json")

        model = ReviewSentimentModel(tokenizer, self.output_dim, False)
        buffer = BytesIO(self.model_state_dict_bytes)
        model.load_state_dict(torch.load(buffer, map_location=device, weights_only=True))

        def predict_sentiment(text, model, tokenizer, device):
            ids = tokenizer(text)["input_ids"]
            tensor = torch.LongTensor(ids).unsqueeze(dim=0).to(device)
            prediction = model(tensor).squeeze(dim=0)
            probability = torch.softmax(prediction, dim=-1)
            predicted_class = prediction.argmax(dim=-1).item()
            predicted_probability = probability[predicted_class].item()
            return predicted_class, predicted_probability

        print("(Clearly these are toy examples; one could load a batch of examples here for more rigorous error analysis)")

        text = "This film is terrible!"
        print(f"Analysis of text: {text}")
        print(predict_sentiment(text, model, tokenizer, device))

        text = "This film is not terrible, it's great!"
        print(f"Analysis of text: {text}")
        print(predict_sentiment(text, model, tokenizer, device))

        text = "This film is not terrible, it's great!"
        print(f"Analysis of text: {text}")
        print(predict_sentiment(text, model, tokenizer, device))

        self.next(self.upload_model_to_gcs)

    @pypi(python='3.10.8',
          packages={
              'mozmlops': '0.1.4'
          })
    @kubernetes
    @step
    def upload_model_to_gcs(self):
        from mozmlops.cloud_storage_api_client import CloudStorageAPIClient

        print(f"Uploading model to gcs")
        # init client
        storage_client = CloudStorageAPIClient(
            project_name=GCS_PROJECT_NAME, bucket_name=GCS_BUCKET_NAME
        )
        storage_client.store(data=self.model_state_dict_bytes, storage_path=MODEL_STORAGE_PATH)
        self.next(self.end)

    @kubernetes
    @step
    def end(self):
        """
        This is the mandatory 'end' step: it prints some helpful information
        to access the model and the used dataset.
        """
        print(
            f"""
            Flow complete.

            """
        )


if __name__ == "__main__":
    ReviewSentimentFlow()