bert-fine-tune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLM Fine-tuning BERT for Text Classification Pipeline Notebook\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/unionai-oss/fine-tune-bert-sentiment-classifier/blob/main/bert-fine-tune.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "GitHub repository: https://github.com/unionai-oss/fine-tune-bert-sentiment-classifier/\n",
    "\n",
    "This notebook is a pipeline for fine-tuning a pre-trained BERT model for text classification using the Hugging Face Trasnforers library.\n",
    "\n",
    "## Project Setup:\n",
    "\n",
    "Sign up for Union while libraries are installing below:\n",
    "\n",
    "- Union sign up: https://signup.union.ai/\n",
    "- Union Platform: https://serverless.union.ai/\n",
    "- Hugging Face Sign up: https://huggingface.co/ (Optional)\n",
    "\n",
    "Install python libraries by running the code cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Install required packages if not already installed (for Colab, etc)\n",
    "# %%bash\n",
    "# pip install union==0.1.82 \\\n",
    "#             transformers==4.39.1 \\\n",
    "#             datasets==2.18.0 \\\n",
    "#             matplotlib==3.8.3 \\\n",
    "#             torch==2.0.1 \\\n",
    "#             accelerate==0.27.2 \\\n",
    "#             scikit-learn==1.5.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install union==0.1.128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Auth Union accont"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union create login --serverless --auth device-flow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a Simple Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile simple_wf.py\n",
    "\n",
    "# Import libraries and modules\n",
    "# task\n",
    "from flytekit import task, workflow\n",
    "\n",
    "@task\n",
    "def hello_world(name: str) -> str:\n",
    "    return f\"Hello {name}\"\n",
    "\n",
    "# workflow\n",
    "@workflow\n",
    "def main(name: str) -> str:\n",
    "    return hello_world(name=name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run locally\n",
    "!union run simple_wf.py main --name Flyte"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run on Union\n",
    "!union run --remote simple_wf.py main --name Flyte"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union run simple_wf.py main --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union run --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine tune BERT for Sentiment Classification Workflow\n",
    "- Create a Container Image\n",
    "- Download Dataset\n",
    "- Download Pre-trained BERT Model\n",
    "- Visualize dataset\n",
    "- Fine-tune BERT Model\n",
    "- Evaluate BERT Model\n",
    "- Save Model to Hugging Face Model Hub\n",
    "- Use Model for Inference on New Data\n",
    "\n",
    "Optional: Get Hugging Face API token (Read & Write Access)\n",
    "\n",
    "https://huggingface.co/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Secrets in Union\n",
    "# Add Hugging Face Secret\n",
    "!union create secret hf_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union get secret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !union delete secret hf_token\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile bert-fine-tune.py\n",
    "\n",
    "# Note: this workflow can be broken into modular python files for better organization.\n",
    "\n",
    "\n",
    "# %% import libraries & Create container image\n",
    "# ---------------------------\n",
    "from pathlib import Path\n",
    "import os\n",
    "from flytekit import (Deck, ImageSpec, Resources, Secret, current_context,\n",
    "                      task, workflow)\n",
    "from flytekit.types.directory import FlyteDirectory\n",
    "from flytekit.types.file import FlyteFile\n",
    "from transformers import BertForSequenceClassification\n",
    "\n",
    "image = ImageSpec(\n",
    "    packages=[\n",
    "        \"union==0.1.128\",\n",
    "        \"flytekit==1.14.0\",\n",
    "        \"transformers==4.39.1\",\n",
    "        \"datasets==2.18.0\",\n",
    "        \"matplotlib==3.8.3\",\n",
    "        \"torch==2.0.1\",\n",
    "        \"accelerate==0.27.2\",\n",
    "        \"scikit-learn==1.5.1\"\n",
    "    ],\n",
    ")\n",
    "\n",
    "# %% download dataset\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    cache=True,\n",
    "    cache_version=\"v3\",\n",
    "    requests=Resources(cpu=\"2\", mem=\"2Gi\"),\n",
    ")\n",
    "def download_dataset() -> FlyteDirectory:\n",
    "    from datasets import load_dataset\n",
    "   \n",
    "    working_dir = Path(current_context().working_directory)\n",
    "    dataset_cache_dir = working_dir / \"dataset_cache\"\n",
    "\n",
    "    load_dataset(\"imdb\", cache_dir=dataset_cache_dir)\n",
    "\n",
    "    return dataset_cache_dir\n",
    "\n",
    "# %% download model\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    cache=True,\n",
    "    cache_version=\"v3\",\n",
    "    requests=Resources(cpu=\"2\", mem=\"2Gi\"),\n",
    ")\n",
    "def download_model(model: str) -> FlyteDirectory:\n",
    "    from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "\n",
    "    working_dir = Path(current_context().working_directory)\n",
    "    model_cache_dir = working_dir / \"model_cache\"\n",
    "\n",
    "    AutoTokenizer.from_pretrained(model, cache_dir=model_cache_dir)\n",
    "    AutoModelForSequenceClassification.from_pretrained(model, cache_dir=model_cache_dir)\n",
    "    return model_cache_dir\n",
    "\n",
    "\n",
    "# %% visualize data\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    enable_deck=True,\n",
    "    requests=Resources(cpu=\"2\", mem=\"2Gi\"),\n",
    ")\n",
    "def visualize_data(dataset_cache_dir: FlyteDirectory):\n",
    "    from datasets import load_dataset\n",
    "    import matplotlib.pyplot as plt\n",
    "    import pandas as pd\n",
    "    import base64\n",
    "    from textwrap import dedent\n",
    "\n",
    "    ctx = current_context()\n",
    "\n",
    "    # Load the dataset\n",
    "    dataset = load_dataset(\"imdb\", cache_dir=dataset_cache_dir)\n",
    "    train_df = pd.DataFrame(dataset['train'])\n",
    "    test_df = pd.DataFrame(dataset['test'])\n",
    "\n",
    "    # Create the deck for visualization\n",
    "    deck = Deck(\"Dataset Analysis\")\n",
    "\n",
    "    # Sample one review from each class (positive and negative) from the training and test datasets\n",
    "    train_positive_review = train_df[train_df['label'] == 1].iloc[0]['text']\n",
    "    train_negative_review = train_df[train_df['label'] == 0].iloc[0]['text']\n",
    "    test_positive_review = test_df[test_df['label'] == 1].iloc[0]['text']\n",
    "    test_negative_review = test_df[test_df['label'] == 0].iloc[0]['text']\n",
    "\n",
    "    # Visualize label distribution for training data\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    train_df['label'].value_counts().plot(kind='bar', color='skyblue')\n",
    "    plt.title('Train Data Label Distribution')\n",
    "    plt.xlabel('Label')\n",
    "    plt.ylabel('Count')\n",
    "    plt.tight_layout()\n",
    "    train_label_dist_path = \"/tmp/train_label_distribution.png\"\n",
    "    plt.savefig(train_label_dist_path)\n",
    "    plt.close()\n",
    "\n",
    "    # Visualize label distribution for test data\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    test_df['label'].value_counts().plot(kind='bar', color='lightgreen')\n",
    "    plt.title('Test Data Label Distribution')\n",
    "    plt.xlabel('Label')\n",
    "    plt.ylabel('Count')\n",
    "    plt.tight_layout()\n",
    "    test_label_dist_path = \"/tmp/test_label_distribution.png\"\n",
    "    plt.savefig(test_label_dist_path)\n",
    "    plt.close()\n",
    "\n",
    "    # Convert images to base64 and embed in HTML\n",
    "    def image_to_base64(image_path):\n",
    "        with open(image_path, \"rb\") as img_file:\n",
    "            return base64.b64encode(img_file.read()).decode(\"utf-8\")\n",
    "\n",
    "    train_image_base64 = image_to_base64(train_label_dist_path)\n",
    "    test_image_base64 = image_to_base64(test_label_dist_path)\n",
    "\n",
    "    # HTML report with styled text, tables, and embedded images\n",
    "    html_report = dedent(f\"\"\"\n",
    "    <div style=\"font-family: Arial, sans-serif; line-height: 1.6;\">\n",
    "        <h2 style=\"color: #2C3E50;\">Dataset Analysis</h2>\n",
    "        \n",
    "        <h3 style=\"color: #2980B9;\">Training Data Summary</h3>\n",
    "        <p>Below is a summary of the training dataset including the distribution of labels.</p>\n",
    "        Shape: {train_df.shape} <br>\n",
    "        Columns: {train_df.columns} <br>\n",
    "        Label Distribution: {train_df['label'].value_counts()} <br>\n",
    "        \n",
    "        <h3 style=\"color: #2980B9;\">Sample Reviews from Training Data</h3>\n",
    "        <p><strong>Positive Review:</strong> {train_positive_review}</p>\n",
    "        <p><strong>Negative Review:</strong> {train_negative_review}</p>\n",
    "\n",
    "        <h3 style=\"color: #2980B9;\">Training Data Label Distribution</h3>\n",
    "        <p>The following bar chart shows the distribution of labels in the training dataset:</p>\n",
    "        <img src=\"data:image/png;base64,{train_image_base64}\" alt=\"Train Data Label Distribution\" width=\"600\">\n",
    "\n",
    "        <h3 style=\"color: #2980B9;\">Test Data Summary</h3>\n",
    "        <p>Below is a summary of the test dataset including the distribution of labels.</p>\n",
    "        Shape: {test_df.shape} <br>\n",
    "        Columns: {test_df.columns} <br>\n",
    "        Label Distribution: {test_df['label'].value_counts()} <br>\n",
    "        \n",
    "        <h3 style=\"color: #2980B9;\">Sample Reviews from Test Data</h3>\n",
    "        <p><strong>Positive Review:</strong> {test_positive_review}</p>\n",
    "        <p><strong>Negative Review:</strong> {test_negative_review}</p>\n",
    "\n",
    "        <h3 style=\"color: #2980B9;\">Test Data Label Distribution</h3>\n",
    "        <p>The following bar chart shows the distribution of labels in the test dataset:</p>\n",
    "        <img src=\"data:image/png;base64,{test_image_base64}\" alt=\"Test Data Label Distribution\" width=\"600\">\n",
    "    </div>\n",
    "    \"\"\")\n",
    "\n",
    "    # Append HTML content to the deck\n",
    "    deck.append(html_report)\n",
    "\n",
    "    # Insert the deck into the context\n",
    "    ctx.decks.insert(0, deck)\n",
    "\n",
    "\n",
    "# %% train model\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    requests=Resources(cpu=\"4\", mem=\"12Gi\", gpu=\"1\"),\n",
    ")\n",
    "def train_model(model_name: str, \n",
    "                dataset_cache_dir: FlyteDirectory,\n",
    "                model_cache_dir: FlyteDirectory,\n",
    "                epochs: int = 3\n",
    "    ) -> BertForSequenceClassification:\n",
    "    from datasets import load_dataset\n",
    "    import numpy as np\n",
    "    from transformers import(\n",
    "        AutoTokenizer,\n",
    "        AutoModelForSequenceClassification,\n",
    "        TrainingArguments,\n",
    "        Trainer,\n",
    "    )\n",
    "    ctx = current_context()\n",
    "\n",
    "    working_dir = Path(ctx.working_directory)\n",
    "    train_dir = working_dir / \"models\"\n",
    "\n",
    "    dataset = load_dataset(\"imdb\", cache_dir=dataset_cache_dir)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=model_cache_dir)\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
    "        model_name,\n",
    "        num_labels=2,\n",
    "        id2label={0: \"NEGATIVE\", 1: \"POSITIVE\"},\n",
    "        label2id={\"NEGATIVE\": 0, \"POSITIVE\": 1},\n",
    "        cache_dir=model_cache_dir,\n",
    "    )\n",
    "\n",
    "    def tokenizer_function(examples):\n",
    "        return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
    "\n",
    "    # Use a small subset such that finetuning completes\n",
    "    small_train_dataset = (\n",
    "        dataset[\"train\"].shuffle(seed=42).select(range(500)).map(tokenizer_function)\n",
    "    )\n",
    "    small_eval_dataset = (\n",
    "        dataset[\"test\"].shuffle(seed=42).select(range(100)).map(tokenizer_function)\n",
    "    )\n",
    "\n",
    "    def compute_metrics(eval_pred):\n",
    "        logits, labels = eval_pred\n",
    "        predictions = np.argmax(logits, axis=-1)\n",
    "        return {\"accuracy\": np.mean(predictions == labels)}\n",
    "    \n",
    "    training_args = TrainingArguments(\n",
    "        output_dir=train_dir,\n",
    "        evaluation_strategy=\"epoch\",\n",
    "        num_train_epochs=epochs,\n",
    "    )\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=small_train_dataset,\n",
    "        eval_dataset=small_eval_dataset,\n",
    "        compute_metrics=compute_metrics,\n",
    "    )\n",
    "    trainer.train()\n",
    "\n",
    "    return model\n",
    "\n",
    "# %% Evaluate model\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    enable_deck=True,\n",
    "    requests=Resources(cpu=\"2\", mem=\"12Gi\", gpu=\"1\"),\n",
    ")\n",
    "def evaluate_model(\n",
    "    model: BertForSequenceClassification,\n",
    "    dataset_cache_dir: FlyteDirectory,\n",
    "    model_cache_dir: FlyteDirectory,\n",
    ") -> dict:\n",
    "    from datasets import load_dataset\n",
    "    from transformers import AutoTokenizer, Trainer, TrainingArguments\n",
    "    from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score\n",
    "    import numpy as np\n",
    "    import torch\n",
    "\n",
    "    # Load the test dataset and tokenizer\n",
    "    dataset = load_dataset(\"imdb\", cache_dir=dataset_cache_dir)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\", cache_dir=model_cache_dir)\n",
    "\n",
    "    def tokenize_function(examples):\n",
    "        return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
    "\n",
    "    # Use a small subset (200 examples) for evaluation\n",
    "    eval_dataset = dataset[\"test\"].shuffle(seed=42).select(range(200)).map(tokenize_function)\n",
    "\n",
    "    def compute_metrics(eval_pred):\n",
    "        logits, labels = eval_pred\n",
    "        predictions = np.argmax(logits, axis=-1)\n",
    "        accuracy = accuracy_score(labels, predictions)\n",
    "        f1 = f1_score(labels, predictions, average=\"weighted\")\n",
    "        precision = precision_score(labels, predictions, average=\"weighted\")\n",
    "        recall = recall_score(labels, predictions, average=\"weighted\")\n",
    "        return {\n",
    "            \"accuracy\": accuracy,\n",
    "            \"f1\": f1,\n",
    "            \"precision\": precision,\n",
    "            \"recall\": recall,\n",
    "        }\n",
    "\n",
    "    # Initialize Trainer for evaluation\n",
    "    training_args = TrainingArguments(\n",
    "        output_dir=\".\", \n",
    "        per_device_eval_batch_size=16, \n",
    "        dataloader_drop_last=False,\n",
    "    )\n",
    "\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        eval_dataset=eval_dataset,\n",
    "        compute_metrics=compute_metrics,\n",
    "    )\n",
    "\n",
    "    # Run evaluation\n",
    "    eval_results = trainer.evaluate()\n",
    "\n",
    "    print(f\"Evaluation results on 100 examples: {eval_results}\")\n",
    "\n",
    "    return eval_results\n",
    "\n",
    "\n",
    "# %% save model\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    requests=Resources(cpu=\"2\", mem=\"2Gi\"),\n",
    "    secret_requests=[Secret(group=None, key=\"hf_token\")],\n",
    ")\n",
    "def save_model(model: BertForSequenceClassification, repo_name: str) -> str:\n",
    "    from huggingface_hub import HfApi\n",
    "\n",
    "    ctx = current_context()\n",
    "    \n",
    "    working_dir = Path(ctx.working_directory)\n",
    "    model_path = working_dir / \"model\"\n",
    "    model.save_pretrained(model_path)\n",
    "\n",
    "    # Ensure the model files are saved\n",
    "    model_files = [\"model.safetensors\", \"config.json\"]\n",
    "    for file_name in model_files:\n",
    "        file_path = model_path / file_name\n",
    "        if not file_path.exists():\n",
    "            raise FileNotFoundError(f\"Model file not found: {file_path}\")\n",
    "    \n",
    "    # dir list for debug\n",
    "    dir_list = os.listdir(model_path)\n",
    "    print(\"Files and directories in '\", model_path, \"' :\")\n",
    "    print(dir_list)\n",
    "\n",
    "\n",
    "     # set hf_token from local or union secret\n",
    "    hf_token = os.getenv(\"HF_TOKEN\")\n",
    "    if hf_token is None:\n",
    "        # If HF_TOKEN is not found, attempt to get it from the Flyte secrets\n",
    "        hf_token = ctx.secrets.get(key=\"hf_token\")\n",
    "        print(\"Using Hugging Face token from Union secrets.\")\n",
    "    else:\n",
    "        print(\"Using Hugging Face token from env.\")\n",
    "\n",
    "    # Create a new repository (if it doesn't exist)\n",
    "    api = HfApi()\n",
    "    api.create_repo(repo_name, token=hf_token, exist_ok=True)\n",
    "\n",
    "    # Upload the model to the HF repository\n",
    "    # Upload each model file to the HF repository\n",
    "    for file_name in model_files:\n",
    "        file_path = model_path / file_name\n",
    "        api.upload_file(\n",
    "            path_or_fileobj=str(file_path),\n",
    "            path_in_repo=file_name,\n",
    "            repo_id=repo_name,\n",
    "            commit_message=f\"Upload {file_name}\",\n",
    "            repo_type=None,\n",
    "            token=hf_token\n",
    "        )\n",
    "    return f\"Model uploaded to Hugging Face Hub: {repo_name}\"\n",
    "\n",
    "# %% Predict sentiment\n",
    "# ---------------------------\n",
    "@task(\n",
    "    container_image=image,\n",
    "    requests=Resources(cpu=\"2\", mem=\"2Gi\", gpu=\"1\"),\n",
    "    retries=2,\n",
    ")\n",
    "def predict_sentiment(model: BertForSequenceClassification, text: str, model_cache_dir: FlyteDirectory) -> dict:\n",
    "    from transformers import AutoTokenizer, pipeline\n",
    "\n",
    "    # Load the tokenizer\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\", cache_dir=model_cache_dir)\n",
    "\n",
    "    # Initialize the pipeline for sentiment analysis\n",
    "    nlp_pipeline = pipeline(\"text-classification\", model=model, tokenizer=tokenizer)\n",
    "\n",
    "    # Perform the prediction\n",
    "    prediction = nlp_pipeline(text)\n",
    "\n",
    "    return prediction[0] \n",
    "\n",
    "# %% Workflow\n",
    "@workflow\n",
    "def bert_ft(model: str = \"bert-base-uncased\",\n",
    "            repo_name: str = \"my-model\",\n",
    "            test_text: str = \"I love this movie!\",\n",
    "            epochs: int = 3\n",
    "            ) -> (BertForSequenceClassification, FlyteDirectory, dict, dict):\n",
    "    \n",
    "    dataset_cache_dir = download_dataset()\n",
    "    model_cache_dir = download_model(model)\n",
    "    visualize_data(dataset_cache_dir=dataset_cache_dir)\n",
    "    model = train_model(model_name=model, \n",
    "                        dataset_cache_dir=dataset_cache_dir, \n",
    "                        model_cache_dir=model_cache_dir,\n",
    "                        epochs=epochs)\n",
    "    eval_results = evaluate_model(model=model, \n",
    "                                  dataset_cache_dir=dataset_cache_dir, \n",
    "                                  model_cache_dir=model_cache_dir)\n",
    "    save_model(model=model, repo_name=repo_name) # uncomment to save model to HF\n",
    "    \n",
    "    # Pass the model and other needed arguments directly to the predict_sentiment task\n",
    "    prediction = predict_sentiment(model=model, text=test_text, model_cache_dir=model_cache_dir)\n",
    "    \n",
    "    # Return results as a dict\n",
    "    return model, model_cache_dir, eval_results, prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the workflow on Union with the `--remote` flag.\n",
    "\n",
    "Replace HFUSERNAME/REPOID with your Hugging Face username and repository ID you want to save the model to. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union run --remote bert-fine-tune.py bert_ft --repo_name HFUSERNAME/REPOID"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try running the workflow with different epochs.  Or relaunch the workflow from the UI with different hyperparameters.\n",
    "\n",
    "```bash\n",
    "\n",
    "!union run --remote bert-fine-tune.py bert_ft --epochs 6 --repo_name HFUSERNAME/REPOID\n",
    "\n",
    "```\n",
    "\n",
    "You can also modify the workflow to include different hyperparameters or models.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union run --remote bert-fine-tune.py bert_ft --epochs 6 --repo_name HFUSERNAME/REPOID\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Model for Inference\n",
    "\n",
    "There are several ways to use the model for inference from here. \n",
    "\n",
    "- Download the model from the Hugging Face Model Hub and use it in your own code.\n",
    "- Deploy the model on a platform like Sagemaker, or Hugging Face spaces that has compute resources.\n",
    "- Use the model directly in Union serverless. Run from API, or UI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use UnionRemote() to run the workflow on Union from the API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Without Union Artifacts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize remote context\n",
    "from union.remote import UnionRemote\n",
    "remote = UnionRemote()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "# search for the 10 most recent executions of the workflow\n",
    "recent_executions = remote.recent_executions(limit=10)\n",
    "executions = [\n",
    "    e for e in recent_executions if e.spec.launch_plan.name == \"bert-fine-tune.bert_ft\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "executions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recent_ex_id = executions[0].id.name\n",
    "print(recent_ex_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "execution = remote.fetch_execution(name=recent_ex_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Flyte Serialized object: Type: <FlyteWorkflowExecution> Value: <id { project: \"default\" domain: \"development\" name: \"fc5519022f1a44c9390c\" } spec { launch_plan { resource_type: LAUNCH_PLAN project: \"default\" domain: \"development\" name: \"bert-fine-tune.bert_ft\" version: \"lSB3Q7nT4OniBMZzxA8aYQ\" } metadata { principal: \"00ugcrzo37k3huEiU5d7\" system_metadata { execution_cluster: \"union-us-central1\" } } disable_all: false labels { } annotations { } auth_role { } } closure { outputs { uri: \"s3://unionai-production-serverless-us-east-2/metadata/sagecodes/default/development/fc5519022f1a44c9390c/offloaded_outputs\" } phase: SUCCEEDED started_at { seconds: 1724885764 nanos: 99438000 } duration { seconds: 545 nanos: 362008000 } created_at { seconds: 1724885762 nanos: 795334000 } updated_at { seconds: 1724886309 nanos: 461446000 } }>]"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "execution\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "execution.outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ft_llm = execution.outputs[\"o0\"].remote_source\n",
    "model_cache = execution.outputs[\"o1\"].remote_source\n",
    "print(ft_llm)\n",
    "print(model_cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fetch a task\n",
    "predict_task = remote.fetch_task(name=\"bert-fine-tune.predict_sentiment\")\n",
    "\n",
    "# Fecth specific version of a task\n",
    "# task = remote.fetch_task(name=\"bert-fine-tune.predict_sentiment\", version=\"_XyPQzsykVFTceNWitQmAg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# view task details\n",
    "predict_task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flytekit.types.file import FlyteFile\n",
    "inputs = {\n",
    "    \"text\": \"Why did this have 40 minutes of farming wheat?\", \n",
    "    \"model_cache_dir\": model_cache,\n",
    "    \"model\": FlyteFile(ft_llm)\n",
    "}\n",
    "\n",
    "# # Execute the task\n",
    "execution = remote.execute(predict_task, inputs=inputs)\n",
    "# execution = remote.execute(task, inputs=inputs, wait=True) # wait for execution to finish\n",
    "\n",
    "url = remote.generate_console_url(execution)\n",
    "print(f\"🚀 Union Serverless execution url: {url}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "synced_execution = remote.sync(execution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synced_execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = synced_execution.outputs[\"o0\"]\n",
    "print(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: If doing something like this for production, I'd recommend:\n",
    "- Using [Union Artifacts](https://www.union.ai/blog-post/data-aware-event-driven-ai-orchestration-with-artifacts) to store the model and dataset\n",
    "- Making inference a separate workflow\n",
    "\n",
    "We're working on new features that will make models more readily available for inference in Union."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "from union.remote import UnionRemote\n",
    "from flytekit.types.file import FlyteFile\n",
    "\n",
    "remote = UnionRemote()\n",
    "\n",
    "# Fetch recent executions\n",
    "recent_executions = remote.recent_executions(limit=10)\n",
    "executions = [\n",
    "    e for e in recent_executions if e.spec.launch_plan.name == \"bert-fine-tune.bert_ft\"\n",
    "]\n",
    "recent_ex_id = executions[0].id.name\n",
    "execution = remote.fetch_execution(name=recent_ex_id)\n",
    "\n",
    "ft_llm = execution.outputs[\"o0\"].remote_source\n",
    "model_cache = execution.outputs[\"o1\"].remote_source\n",
    "\n",
    "predict_task = remote.fetch_task(name=\"bert-fine-tune.predict_sentiment\")\n",
    "\n",
    "\n",
    "def execute_flyte_task(text):\n",
    "    inputs = {\n",
    "        \"text\": text,\n",
    "        \"model_cache_dir\": model_cache,\n",
    "        \"model\": FlyteFile(ft_llm)\n",
    "    }\n",
    "    execution = remote.execute(predict_task, inputs=inputs, wait=True)\n",
    "\n",
    "    # Extract the response\n",
    "    response = execution.outputs[\"o0\"]\n",
    "\n",
    "    # Format the response\n",
    "    sentiment = response['label']\n",
    "    score = response['score']\n",
    "\n",
    "    # Determine color based on sentiment\n",
    "    color = \"red\" if sentiment == \"NEGATIVE\" else \"green\"\n",
    "\n",
    "    # Format the output HTML\n",
    "    output_html = f\"\"\"\n",
    "    <div style=\"text-align: center;\">\n",
    "        <h2>Sentiment: <span style=\"color: {color};\">{sentiment}</span></h2>\n",
    "        <p>Confidence Score: {score:.2f}</p>\n",
    "    </div>\n",
    "    \"\"\"\n",
    "\n",
    "    return output_html\n",
    "\n",
    "\n",
    "# Launch Gradio app\n",
    "iface = gr.Interface(\n",
    "    fn=execute_flyte_task,\n",
    "    inputs=[\"text\"],\n",
    "    outputs=gr.HTML(),  # Change output to HTML for better formatting\n",
    "    live=False,\n",
    ")\n",
    "\n",
    "iface.launch(debug=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!union create app gradio-app"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Resources:\n",
    "- [RAG on Union](https://github.com/unionai-oss/union-rag) - [Video tutorial](https://www.youtube.com/watch?v=IpAa1_oHNWs)\n",
    "- [a100s on Union Serverless ](https://www.union.ai/blog-post/union-serverless-broadens-support-for-nvidia-accelerated-computing)\n",
    "- [How LinkedIn uses Flyte for all LLM workflows](https://www.youtube.com/live/kl-fMSe2_l8?si=hEBV-LYB_SrY_1Au&t=217)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bert",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}