"Traceback (most recent call last):\n\n      File \"/usr/local/lib/python3.10/dist-packages/flytekit/exceptions/scopes.py\", line 219, in user_entry_point\n        return wrapped(*args, **kwargs)\n      File \"/root/workflows/hf_example.py\", line 58, in wf\n        train_dataset, eval_dataset = get_and_tokenize_datasets()\n\nMessage:\n\n    cannot unpack non-iterable VoidPromise object\n\nUser error."}

_import_ logging

_import_ os

_from_ typing _import_ Tuple, Dict, List

_from_ flytekit _import_ current_context, task, Resources, workflow

_from_ datasets _import_ load_dataset, Dataset

_from_ transformers _import_ AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments

_from_ transformers.integrations _import_ FlyteCallback

_import_ numpy _as_ np

_import_ evaluate

logger = logging.getLogger(__file__)

model_name = "bert-base-cased"

tokenizer = AutoTokenizer.from_pretrained(model_name)

metric = evaluate.load("accuracy")

def tokenize_function(examples: Dict[str, List[int]]) -> Dict[str, List[List[int]]]:

_return_ tokenizer(examples["text"], padding="max_length", truncation=True)

def compute_metrics(eval_pred):

logits, labels = eval_pred

predictions = np.argmax(logits, axis=-1)

_return_ metric.compute(predictions=predictions, references=labels)

@task

def get_and_tokenize_datasets() -> Tuple[Dataset, Dataset]:

dataset = load_dataset("yelp_review_full")

tokenized_datasets = dataset.map(tokenize_function, batched=True)

small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))

small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

_return_ small_train_dataset, small_eval_dataset

@task(requests=Resources(gpu="1"), disable_deck=False, interruptible=False)

def train_hf_transformer(train_dataset: Dataset, eval_dataset: Dataset) -> str:

model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=5)

cp = current_context().checkpoint

training_args = TrainingArguments(

"trainer",

save_strategy="epoch",

num_train_epochs=5,

evaluation_strategy="epoch",

train_dataset=train_dataset,

eval_dataset=eval_dataset,

compute_metrics=compute_metrics,

)

trainer = Trainer(model,training_args, callbacks=[FlyteCallback()])

trainer.train(resume_from_checkpoint=cp.restore())

path = os.path.join(current_context().working_directory, "example_hf_flyte_model")

trainer.save_model(path)

_return_ path

@workflow

def wf() -> str:

train_dataset, eval_dataset = get_and_tokenize_datasets()

path = train_hf_transformer(train_dataset=train_dataset, eval_dataset=eval_dataset)

_return_ path

_if_ __name__ == "__main__":

wf()

FROM <http://nvcr.io/nvidia/pytorch:23.10-py3|nvcr.io/nvidia/pytorch:23.10-py3>

WORKDIR /root

ENV LANG C.UTF-8

ENV LC_ALL C.UTF-8

ENV PYTHONPATH /root

COPY requirements.txt /root

RUN pip install -r /root/requirements.txt

COPY . /root

# This tag is supplied by the build script and will be used to determine the version

# when registering tasks, workflows, and launch plans

ARG tag

ENV FLYTE_INTERNAL_IMAGE $tag

pyflyte --pkgs workflows package

flytectl register files

WARNING  {"asctime": "2024-01-11 16:41:05,397", "name": "flytekit", "levelname": "WARNING", "message": "Unsupported Type <class 'datasets.arrow_dataset.Dataset'> found, Flyte will default to use PickleFile as the transport. Pickle can only be used to send objects between the exact same version of Python, and we strongly recommend to use python type that flyte support."}

import logging
import os
from typing import Tuple, Dict, List
from flytekit import current_context, task, Resources, workflow
from datasets import load_dataset, Dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from transformers.integrations import FlyteCallback
import numpy as np
import evaluate

logger = logging.getLogger(__file__)
model_name = "bert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
metric = evaluate.load("accuracy")


def tokenize_function(examples: Dict[str, List[int]]) -> Dict[str, List[List[int]]]:
    return tokenizer(examples["text"], padding="max_length", truncation=True)


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)


@task
def get_and_tokenize_datasets() -> Tuple[Dataset, Dataset]:
    dataset = load_dataset("yelp_review_full")
    tokenized_datasets = dataset.map(tokenize_function, batched=True)
    small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
    small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))
    return small_train_dataset, small_eval_dataset


@task(requests=Resources(gpu="1"), disable_deck=False, interruptible=False)
def train_hf_transformer(train_dataset: Dataset, eval_dataset: Dataset) -> str:
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=5)
    cp = current_context().checkpoint
    training_args = TrainingArguments(
        "trainer",
        save_strategy="epoch",
        num_train_epochs=5,
        evaluation_strategy="epoch",
    )
    trainer = Trainer(
        model,
        training_args,
        callbacks=[FlyteCallback()],
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
    )
    trainer.train(resume_from_checkpoint=cp.restore())
    path = os.path.join(current_context().working_directory, "example_hf_flyte_model")
    trainer.save_model(path)
    return path


@workflow
def wf() -> str:
    train_dataset, eval_dataset = get_and_tokenize_datasets()
    path = train_hf_transformer(train_dataset=train_dataset, eval_dataset=eval_dataset)
    return path


if __name__ == "__main__":
    wf()

pyflyte run --remote

pyflyte register

remote.execute