zero-data-model-foundry/zero-data-model-foundry.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Building A Zero-Data Model Foundry\n",
        "\n",
        "## Setting Up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "import time\n",
        "import json\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "from dataclasses import dataclass\n",
        "from pathlib import Path\n",
        "from uuid import uuid4\n",
        "\n",
        "import covalent as ct\n",
        "import covalent_cloud as cc\n",
        "import torch\n",
        "from covalent_cloud.cloud_executor.models.gpu import GPU_TYPE\n",
        "from datasets import Dataset, load_from_disk\n",
        "from peft import LoraConfig\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    AutoTokenizer,\n",
        "    BitsAndBytesConfig,\n",
        "    TrainingArguments,\n",
        "    pipeline\n",
        ")\n",
        "from trl import SFTTrainer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Authenticating with Covalent Cloud"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "CC_API_KEY = os.environ[\"CC_API_KEY\"]  # set in `environment.yml` file\n",
        "cc.save_api_key(CC_API_KEY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Create a [cloud volume](https://docs.covalent.xyz/docs/cloud/guides/cloud_storage) for persistent storage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "volume = cc.volume(\"model-storage\")  # store fine-tuned models and generated datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Create [runtime environments](https://docs.covalent.xyz/docs/cloud/guides/cloud_custom_environments) for tasks and services"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "An environment for fine-tuning models:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Environment Already Exists.\n"
          ]
        }
      ],
      "source": [
        "FT_ENV = \"model-fine-tuning\"  # assign unique name for referring to this env\n",
        "\n",
        "cc.create_env(\n",
        "    name=FT_ENV,\n",
        "    pip=[\n",
        "        \"accelerate==0.29.1\",\n",
        "        \"bitsandbytes==0.43.0\",\n",
        "        \"datasets==2.18.0\",\n",
        "        \"pandas==2.2.1\",\n",
        "        \"scipy==1.12.0\",\n",
        "        \"sentencepiece==0.2.0\",\n",
        "        \"torch==2.2.2\",\n",
        "        \"transformers==4.39.3\",\n",
        "        \"trl==0.8.1\",\n",
        "        \"tqdm==4.66.2\",\n",
        "        \"peft==0.10.0\",\n",
        "    ],\n",
        "    wait=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Another environment for running the data generator LLM:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Environment Already Exists.\n"
          ]
        }
      ],
      "source": [
        "DATA_ENV = \"data-generation\"\n",
        "cc.create_env(name=DATA_ENV, pip=[\"vllm\"], wait=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Service: Data Generator LLM\n",
        "\n",
        "This service hosts a powerful LLM that generates synthetic data for fine tuning another model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"./assets/data-generator.png\" alt=\"Highlight data generator component\" height=550px/>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Backend for the Data Generator LLM\n",
        "\n",
        "- H100 GPU\n",
        "\n",
        "- 48 GB RAM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "data_generator_ex = cc.CloudExecutor(\n",
        "    env=DATA_ENV,\n",
        "    num_cpus=6,\n",
        "    num_gpus=1,\n",
        "    gpu_type=GPU_TYPE.H100,\n",
        "    memory=\"48GB\",\n",
        "    time_limit=\"12 hours\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "PROMPT_TEMPLATE = (\n",
        "    \"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\"\n",
        "    \"You are a knowledgeable assistant who generates fine-tuning data for an LLM. \"\n",
        "    \"Please generate {target_items_per_response} data items for the fine-tuning task specified by the user.\\n\"\n",
        "    \"IMPORTANT: Return a JSON array of new items in the format: \\\"{return_format}\\\"\"\n",
        "    \"<|eot_id|>\"\n",
        "    \"<|start_header_id|>user<|end_header_id|>{user_prompt}<|eot_id|>\"\n",
        "    \"<|start_header_id|>assistant<|end_header_id|>\"\n",
        ")\n",
        "\n",
        "\n",
        "@cc.service(executor=data_generator_ex, name=\"LLM Data Generator\", auth=False)\n",
        "def llm_data_generator(model_name):\n",
        "\n",
        "    \"\"\"Uses a powerful LLM to generate synthetic training data.\"\"\"\n",
        "\n",
        "    from vllm import LLM\n",
        "\n",
        "    return {\n",
        "        \"llm\": LLM(model=model_name, trust_remote_code=True, enforce_eager=True),\n",
        "    }\n",
        "\n",
        "@llm_data_generator.endpoint(\"/generate-data\")\n",
        "def generate_data(\n",
        "    llm,\n",
        "    task,\n",
        "    return_format=\"[item] ## [label]\",\n",
        "    num_generations=5,\n",
        "    target_items_per_response=5,\n",
        "):\n",
        "    \"\"\"Generate data for a specific task.\"\"\"\n",
        "\n",
        "    from vllm import SamplingParams\n",
        "\n",
        "    # Format task into prompt.\n",
        "    prompt_ = PROMPT_TEMPLATE.format(\n",
        "        target_items_per_response=target_items_per_response,\n",
        "        return_format=return_format,\n",
        "        user_prompt=json.dumps({\n",
        "            \"task\": task,\n",
        "            \"constraint\": \"Respond with ONLY the generated data as a valid JSON array!\",\n",
        "        }),\n",
        "    )\n",
        "\n",
        "    def _seeded_sampling_params():\n",
        "        seed = random.randint(0, 1_000_000)\n",
        "        return SamplingParams(temperature=0.9, top_p=0.8, max_tokens=2000, seed=seed)\n",
        "\n",
        "    # Create a batch of prompts and sampling params\n",
        "    prompts_batch = [prompt_] * num_generations\n",
        "    params_batch = [_seeded_sampling_params() for _ in range(num_generations)]\n",
        "\n",
        "    # Generate data\n",
        "    outputs = llm.generate(prompts_batch, params_batch)\n",
        "\n",
        "    # Extract and filter generated data\n",
        "    data_items = []\n",
        "    for output in outputs:\n",
        "        generated_text = output.outputs[0].text\n",
        "        try:\n",
        "            data_items.extend(json.loads(generated_text))\n",
        "        except Exception:\n",
        "            continue\n",
        "\n",
        "    return data_items"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Service: New Fine-tuned Models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"./assets/finetune-service.png\" alt=\"Highlight fine-tuned model component\" height=550px/>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Backend for the Fine-tuned Model Service\n",
        "\n",
        "Run on:\n",
        "- L40 GPU\n",
        "- 48 GB RAM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "ft_service_ex = cc.CloudExecutor(\n",
        "    env=FT_ENV,\n",
        "    num_cpus=25,\n",
        "    num_gpus=1,\n",
        "    gpu_type=GPU_TYPE.L40,\n",
        "    memory=\"48GB\",\n",
        "    time_limit=\"12 hours\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "@cc.service(executor=ft_service_ex, volume=volume)\n",
        "def finetuned_llm_service(ft_model_path):\n",
        "\n",
        "    \"\"\"Hosts a fine-tuned LLM for text generation.\"\"\"\n",
        "\n",
        "    ft_model_path_ = Path(\"/tmp\") / Path(ft_model_path).name\n",
        "    if ft_model_path_.exists():\n",
        "        shutil.rmtree(ft_model_path_)\n",
        "    shutil.copytree(ft_model_path, ft_model_path_)\n",
        "\n",
        "    # Load and configure saved model\n",
        "    model = AutoModelForCausalLM.from_pretrained(ft_model_path_, device_map=\"auto\", do_sample=True)\n",
        "\n",
        "    # Load and configure tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(ft_model_path_)\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "    tokenizer.padding_side = \"right\"\n",
        "\n",
        "    # Combine model and tokenizer into a pipeline\n",
        "    pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
        "\n",
        "    return {\"pipe\": pipe, \"model\": model, \"tokenizer\": tokenizer}\n",
        "\n",
        "\n",
        "@finetuned_llm_service.endpoint(\"/generate\")\n",
        "def generate_text(pipe, prompt, max_new_tokens=300):\n",
        "\n",
        "    \"\"\"Generate text from a prompt using the fine-tuned language model.\"\"\"\n",
        "\n",
        "    output = pipe(prompt, truncation=True, max_new_tokens=max_new_tokens, num_return_sequences=1)\n",
        "    return output[0][\"generated_text\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Service: Main Agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"./assets/main-agent.png\" alt=\"Highlight main agent component\" height=550px/>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "agent_ex = cc.CloudExecutor(env=FT_ENV, num_cpus=2, memory=\"12GB\", time_limit=\"12 hours\")\n",
        "\n",
        "@cc.service(executor=agent_ex, name=\"Fine Tuning Launcher\", auth=False, volume=volume)\n",
        "def fine_tuning_launcher(lattice, data_generator_api):\n",
        "\n",
        "    \"\"\"Launches the fine-tuning workflow and deploys the fine-tuned model.\"\"\"\n",
        "\n",
        "    return {\"finetune_lattice\": lattice, \"data_generator_api\": data_generator_api}\n",
        "\n",
        "\n",
        "@fine_tuning_launcher.endpoint(\"/submit\", streaming=True)\n",
        "def submit(\n",
        "    finetune_lattice, data_generator_api,\n",
        "    *,\n",
        "    task,\n",
        "    name,\n",
        "    data_format=\"[item] ## [label]\",\n",
        "    num_generations=5,\n",
        "    target_items_per_response=10,\n",
        "    min_new_examples=2000,\n",
        "    model_to_finetune=\"TinyLlama/TinyLlama-1.1B-Chat-v1.0\",\n",
        "):\n",
        "    \"\"\"Receives a task description, generates fine-tuning data,\n",
        "    and dispatches the fine-tuning + deployment workflow.\"\"\"\n",
        "\n",
        "    yield \"Generating fine-tuning data \"\n",
        "\n",
        "    iteration = 1\n",
        "    new_examples = []\n",
        "    while len(new_examples) < min_new_examples:\n",
        "\n",
        "        texts = data_generator_api.generate_data(\n",
        "            task=task,\n",
        "            return_format=data_format,\n",
        "            num_generations=num_generations,\n",
        "            target_items_per_response=target_items_per_response,\n",
        "        )\n",
        "        new_examples.extend(texts)\n",
        "\n",
        "        yield \".\"\n",
        "        iteration += 1\n",
        "\n",
        "    yield f\"\\nGenerated {len(new_examples)} total examples\\n\"\n",
        "\n",
        "    dataset_save_path = volume / f\"data_{len(new_examples)}-{uuid4()}\"\n",
        "    yield f\"Saving dataset at {dataset_save_path!s}\\n\"\n",
        "    dataset = Dataset.from_dict({\"text\": new_examples})\n",
        "    dataset_save_path.mkdir(parents=True, exist_ok=True)\n",
        "    dataset.save_to_disk(dataset_save_path)\n",
        "\n",
        "    cc.save_api_key(CC_API_KEY)\n",
        "\n",
        "    yield \"\\nDispatching fine-tuning workflow\\n\"\n",
        "    finetuned_llm_service.name = name\n",
        "    dispatch_id = cc.dispatch(finetune_lattice, volume=volume)(\n",
        "        model_to_finetune, str(dataset_save_path), finetuned_llm_service\n",
        "    )\n",
        "    yield f\"Dispatch ID:\\n{dispatch_id}\\n\"\n",
        "    yield \"Fine tuning new model \"\n",
        "    result = None\n",
        "    while result is None:\n",
        "        res = cc.get_result(dispatch_id)\n",
        "        res.result.load()\n",
        "        result = res.result.value\n",
        "        time.sleep(10)\n",
        "        yield \".\"\n",
        "\n",
        "    yield f\"\\nNew Service ID: {result.function_id!s}\\n\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Workflow: Fine Tune & Deploy\n",
        "\n",
        "This workflow runs model fine-tuning on a powerful GPU and deploys the model as a service."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"./assets/finetune-workflow.png\" alt=\"Highlight fine-tune and deploy workflow\" height=550px/>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Training configuration params\n",
        "\n",
        "This dataclass holds the myriad fine-tuning parameter defaults for the PEFT/LoRA approach."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [],
      "source": [
        "@dataclass\n",
        "class FineTuneArguments:\n",
        "    # BitAndBytesConfig\n",
        "    load_in_4bit: bool = True\n",
        "    bnb_4bit_quant_type: str = \"nf4\"\n",
        "    bnb_4bit_compute_dtype: str = \"float16\"\n",
        "    bnb_4bit_use_double_quant: bool = False\n",
        "\n",
        "    # TrainingArguments\n",
        "    output_dir: str = \"./outputs\"\n",
        "    learning_rate: float = 2e-3\n",
        "    num_train_epochs: int = 5\n",
        "    save_total_limit: int = 1\n",
        "    save_strategy: str = \"epoch\"\n",
        "    per_device_train_batch_size: int = 2\n",
        "    gradient_accumulation_steps: int = 1\n",
        "    optim: str = \"paged_adamw_32bit\"\n",
        "    weight_decay: float = 0.001\n",
        "    fp16: bool = False\n",
        "    bf16: bool = False\n",
        "    max_grad_norm: float = 0.3\n",
        "    max_steps: int = -1\n",
        "    warmup_ratio: float = 0.03\n",
        "    group_by_length: bool = True\n",
        "    lr_scheduler_type: str = \"cosine\"\n",
        "    report_to: str = \"none\"\n",
        "\n",
        "    # LoraConfig\n",
        "    lora_alpha: int = 32\n",
        "    lora_dropout: float = 0.05\n",
        "    r: int = 32\n",
        "    bias: str = \"none\"\n",
        "    task_type: str = \"CAUSAL_LM\"\n",
        "\n",
        "    # SFTTrainer\n",
        "    dataset_text_field: str = \"text\"\n",
        "    max_seq_length: int = 1024\n",
        "    packing: bool = True\n",
        "    dataset_batch_size: int = 10\n",
        "\n",
        "    @property\n",
        "    def training_args(self):\n",
        "        return TrainingArguments(\n",
        "            output_dir=self.output_dir,\n",
        "            num_train_epochs=self.num_train_epochs,\n",
        "            per_device_train_batch_size=self.per_device_train_batch_size,\n",
        "            gradient_accumulation_steps=self.gradient_accumulation_steps,\n",
        "            optim=self.optim,\n",
        "            save_strategy=self.save_strategy,\n",
        "            save_total_limit=self.save_total_limit,\n",
        "            learning_rate=self.learning_rate,\n",
        "            weight_decay=self.weight_decay,\n",
        "            fp16=self.fp16,\n",
        "            bf16=self.bf16,\n",
        "            max_grad_norm=self.max_grad_norm,\n",
        "            max_steps=self.max_steps,\n",
        "            warmup_ratio=self.warmup_ratio,\n",
        "            group_by_length=self.group_by_length,\n",
        "            lr_scheduler_type=self.lr_scheduler_type,\n",
        "            report_to=self.report_to,\n",
        "        )\n",
        "\n",
        "    @property\n",
        "    def lora_config(self):\n",
        "        return LoraConfig(\n",
        "            lora_alpha=self.lora_alpha,\n",
        "            lora_dropout=self.lora_dropout,\n",
        "            r=self.r,\n",
        "            bias=self.bias,\n",
        "            task_type=self.task_type,\n",
        "        )\n",
        "\n",
        "    @property\n",
        "    def trainer_params(self):\n",
        "        return {\n",
        "            \"dataset_text_field\": self.dataset_text_field,\n",
        "            \"max_seq_length\": self.max_seq_length,\n",
        "            \"packing\": self.packing,\n",
        "            \"dataset_batch_size\": self.dataset_batch_size,\n",
        "        }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Electrons (i.e. workflow tasks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Fine-tuning task"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Run on:\n",
        "- A100 GPU\n",
        "- 32 GB RAM\n",
        "\n",
        "Tasks exit and release resources after completion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [
        "fine_tune_ex = cc.CloudExecutor(\n",
        "    env=FT_ENV,\n",
        "    num_cpus=6,\n",
        "    num_gpus=1,\n",
        "    gpu_type=GPU_TYPE.A100,\n",
        "    memory=\"32GB\",\n",
        "    time_limit=\"6 hours\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [],
      "source": [
        "@ct.electron(executor=fine_tune_ex)\n",
        "def fine_tune_model_peft(model_path, dataset_path):\n",
        "\n",
        "    \"\"\"Run fine-tuning, save the model, and return the path to the saved model.\"\"\"\n",
        "\n",
        "    ft_args = FineTuneArguments()\n",
        "\n",
        "    # Quantization configuration\n",
        "    quant_config = BitsAndBytesConfig(\n",
        "        load_in_4bit=ft_args.load_in_4bit,\n",
        "        bnb_4bit_quant_type=ft_args.bnb_4bit_quant_type,\n",
        "        bnb_4bit_compute_dtype=getattr(torch, ft_args.bnb_4bit_compute_dtype),\n",
        "        bnb_4bit_use_double_quant=ft_args.bnb_4bit_use_double_quant,\n",
        "    )\n",
        "\n",
        "    # Load and configure the downloaded model from pretrained\n",
        "    model = AutoModelForCausalLM.from_pretrained(\n",
        "        model_path,\n",
        "        quantization_config=quant_config,\n",
        "        device_map=\"auto\",\n",
        "        do_sample=True,\n",
        "    )\n",
        "    model.config.use_cache = False\n",
        "    model.config.pretraining_tp = 1\n",
        "\n",
        "    # Load and configure the tokenizer\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "    tokenizer.padding_side = \"right\"\n",
        "\n",
        "    # Load dataset\n",
        "    dataset_path_ = Path(\"/tmp\") / Path(dataset_path).name\n",
        "    shutil.copytree(dataset_path, dataset_path_)\n",
        "    dataset_path = dataset_path_\n",
        "    dataset = load_from_disk(dataset_path, keep_in_memory=True)\n",
        "\n",
        "    # Set up supervised fine-tuning trainer\n",
        "    trainer = SFTTrainer(\n",
        "        model=model,\n",
        "        train_dataset=dataset,\n",
        "        peft_config=ft_args.lora_config,\n",
        "        tokenizer=tokenizer,\n",
        "        args=ft_args.training_args,\n",
        "        **ft_args.trainer_params,\n",
        "    )\n",
        "\n",
        "    # Run training\n",
        "    trainer.train()\n",
        "\n",
        "    # Save trained model\n",
        "    new_model_path = volume / (model_path.split(\"/\")[-1] + f\"_{uuid4()}\")\n",
        "    trainer.model.save_pretrained(new_model_path)\n",
        "    trainer.tokenizer.save_pretrained(new_model_path)\n",
        "\n",
        "    return new_model_path"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Workflow: Fine tuning (launched by Main Agent)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [],
      "source": [
        "cpu_ex = cc.CloudExecutor(env=FT_ENV, num_cpus=2, memory=\"12GB\")\n",
        "\n",
        "@ct.lattice(executor=cpu_ex, workflow_executor=cpu_ex)\n",
        "def finetune_workflow(model_id, data_path, llm_service):\n",
        "\n",
        "    \"\"\"Run fine tuning, then deploy the fine tuned model.\"\"\"\n",
        "\n",
        "    ft_model_path = fine_tune_model_peft(model_id, data_path)\n",
        "    service_info = llm_service(ft_model_path)\n",
        "\n",
        "    return service_info"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Workflow: Setting up The Zero-Data Model Foundry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [],
      "source": [
        "@ct.lattice(executor=cpu_ex, workflow_executor=cpu_ex)\n",
        "def ai_foundry_setup_workflow(ft_workflow, data_generator_model=\"unsloth/llama-3-8b-Instruct\"):\n",
        "\n",
        "    \"\"\"Set up the data generator LLM and the launcher service.\"\"\"\n",
        "\n",
        "    data_generator_handle = llm_data_generator(data_generator_model)\n",
        "    main_agent = fine_tuning_launcher(ft_workflow, data_generator_handle)\n",
        "    return main_agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ee4e13d5969a4068ad61f3689937ba29",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Output()"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
            ],
            "text/plain": []
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Workflow Dispatch ID:  d9938b2b-246a-4a1e-b7ec-2fbd652f8ce3\n",
            "Main Agent Service ID:  66fc01be50e1b9bb8cb4cb8b\n"
          ]
        }
      ],
      "source": [
        "dispatch_id = cc.dispatch(ai_foundry_setup_workflow, volume=volume)(ft_workflow=finetune_workflow)\n",
        "\n",
        "print(\"Workflow Dispatch ID: \", dispatch_id)\n",
        "\n",
        "res = cc.get_result(dispatch_id, wait=True)\n",
        "res.result.load()\n",
        "main_agent = res.result.value\n",
        "\n",
        "print(\"Main Agent Service ID: \", main_agent.function_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Invoking the Agent API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
            ],
            "text/plain": []
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "╭────────────────────────────── Deployment Information ───────────────────────────────╮\n",
            "│  Name          Fine Tuning Launcher                                                 │\n",
            "│  Description   Launches the fine-tuning workflow and deploys the fine-tuned model.  │\n",
            "│  Function ID   66fc01be50e1b9bb8cb4cb8b                                             │\n",
            "│  Address       https://fn-a.prod.covalent.xyz/66fc01be50e1b9bb8cb4cb8b              │\n",
            "│  Status        ACTIVE                                                               │\n",
            "│  Tags                                                                               │\n",
            "│  Auth Enabled  No                                                                   │\n",
            "╰─────────────────────────────────────────────────────────────────────────────────────╯\n",
            "╭──────────────────────────────────────────────────────────────────────────╮\n",
            "│ \u001b[3m                              POST /submit                              \u001b[0m │\n",
            "│  Streaming    Yes                                                        │\n",
            "│  Description  Receives a task description, generates fine-tuning data,   │\n",
            "│                   and dispatches the fine-tuning + deployment workflow.  │\n",
            "╰──────────────────────────────────────────────────────────────────────────╯\n",
            "\n"
          ]
        }
      ],
      "source": [
        "main_agent = cc.get_deployment(main_agent.function_id)\n",
        "print(main_agent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Example 1: Spoiler Detector"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating fine-tuning data ........................\n",
            "Generated 1000 total examples\n",
            "Saving dataset at /volumes/model-storage/data_1000-efa18179-3651-411d-ba16-c39b374e4d6e\n",
            "\n",
            "Dispatching fine-tuning workflow\n",
            "Dispatch ID:\n",
            "63364307-0602-43a0-9fc8-2ef78c6ba874\n",
            "Fine tuning new model .......................................................................................................................\n",
            "New Service ID: 66fc14ad50e1b9bb8cb4cb8e\n"
          ]
        }
      ],
      "source": [
        "task = \"Fine-tuning an LLM to detect whether or not a movie review contains a spoiler.\"\n",
        "for k in main_agent.submit(\n",
        "    task=task,\n",
        "    name=\"Fine-Tuned Spoiler Detector\",\n",
        "    min_new_examples=1000\n",
        "):\n",
        "    print(k.decode(), end=\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Load a client for the new spoiler agent\n",
        "spoiler_agent = cc.get_deployment(\"66fc14ad50e1b9bb8cb4cb8e\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
            ],
            "text/plain": []
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "╭──────────────────────── Deployment Information ─────────────────────────╮\n",
            "│  Name          Fine-Tuned Spoiler Detector                              │\n",
            "│  Description   Hosts a fine-tuned LLM for text generation.              │\n",
            "│  Function ID   66fc14ad50e1b9bb8cb4cb8e                                 │\n",
            "│  Address       https://fn-c.prod.covalent.xyz/66fc14ad50e1b9bb8cb4cb8e  │\n",
            "│  Status        ACTIVE                                                   │\n",
            "│  Tags                                                                   │\n",
            "│  Auth Enabled  Yes                                                      │\n",
            "╰─────────────────────────────────────────────────────────────────────────╯\n",
            "╭─────────────────────────────────────────────────────────────────────────────────╮\n",
            "│ \u001b[3m                                POST /generate                                 \u001b[0m │\n",
            "│  Streaming    No                                                                │\n",
            "│  Description  Generate text from a prompt using the fine-tuned language model.  │\n",
            "╰─────────────────────────────────────────────────────────────────────────────────╯\n",
            "Authorization token: HdEUB7rjI4cNlak_06XH34ebqrS8iizsVw5eVs0mSja97f0egiq2RrxBxgj7mcKHn1IghSMZRFgCinfQ2tIjRQ\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(spoiler_agent)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'This ripping action-adventure features stellar effects and a superb lead performance from Owen Teague as a timid simian who must rescue his clan from the clutches of a warlike tribe. ## NOT A SPOILER'"
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spoiler_agent.generate(\n",
        "    prompt=(\n",
        "        \"This ripping action-adventure features stellar effects and a \"\n",
        "        \"superb lead performance from Owen Teague as a timid simian \"\n",
        "        \"who must rescue his clan from the clutches of a warlike tribe. ##\"\n",
        "    )\n",
        ")\n",
        "# Review of Kingdom of the Planet of the Apes (2024). Scored 90/100. [source: Metacritic]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'No Spoiler'"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spoiler_agent.generate(prompt=\"Add a lot of dull acting -- except Sir Ian McKellen and Andy Serkis -- and you have an uneven movie with yawns aplenty. ##\").split(\"##\")[-1].strip()\n",
        "# Review of The Lord of the Rings: The Return of the King (2003). Scored 0/100. [source: Metacritic]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'spoiler'"
            ]
          },
          "execution_count": 23,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spoiler_agent.generate(prompt=\"In the end, turns out Soylent green was people! Gross! ##\").split(\"##\")[-1].strip()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'spoiler'"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spoiler_agent.generate(prompt=\"So the Planet of The Apes was Earth all along. That's a lame twist. ##\").split(\"##\")[-1].strip()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Example 2: Grammar Corrector"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating fine-tuning data ........................................\n",
            "Generated 2010 total examples\n",
            "Saving dataset at /volumes/model-storage/data_2010-ae48f9c4-358c-4e56-9373-613fee904878\n",
            "\n",
            "Dispatching fine-tuning workflow\n",
            "Dispatch ID:\n",
            "5af40b24-52fe-48c2-9e91-a6febed0e14e\n",
            "Fine tuning new model ......................................................................................................................................\n",
            "New Service ID: 66fc1d7750e1b9bb8cb4cb91\n"
          ]
        }
      ],
      "source": [
        "task = \"Examples of bot responses that correct grammatical errors.\"\n",
        "data_format = '\"<|user|>{input_sentence}</s><|assistant|>{corrected_sentence}\"'\n",
        "\n",
        "for k in main_agent.submit(\n",
        "    task=task,\n",
        "    name=\"Fine-Tuned Grammar Corrector\",\n",
        "    data_format=data_format\n",
        "):\n",
        "    print(k.decode(), end=\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {},
      "outputs": [],
      "source": [
        "grammar_agent = cc.get_deployment(\"66fc1d7750e1b9bb8cb4cb91\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "I should have never gotten a pet.\n",
            "Jerry and I argued about it.\n",
            "He said, 'No cat bites its own tail.'\n",
            "If I had to choose, I would rather get a dog than a cat.\n",
            "All dogs bite their own tails.\n"
          ]
        }
      ],
      "source": [
        "prompt = \"<|user|>{}</s><|assistant|>\"\n",
        "def correct_grammar(sentence):\n",
        "    prompt_ = prompt.format(sentence)\n",
        "    response = grammar_agent.generate(prompt=prompt_).split(\"<|assistant|>\")[-1].strip()\n",
        "    print(response)\n",
        "\n",
        "correct_grammar(\"I should of never got a pet.\")  # should've\n",
        "correct_grammar(\"Jerry and me argued about it.\")  # Jerry and I\n",
        "correct_grammar(\"He said, 'No cat bites it's own tail.'\")  # its\n",
        "correct_grammar(\"But if I had to choose, Id rather get a dog then a cat.\")  # I'd, than\n",
        "correct_grammar(\"All dogs bite they're own tails.\")  # their"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Example 3: Emoji translation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating fine-tuning data ....................................................\n",
            "Generated 2040 total examples\n",
            "Saving dataset at /volumes/model-storage/data_2040-d173489b-95e5-4752-99e6-42e21435ad07\n",
            "\n",
            "Dispatching fine-tuning workflow\n",
            "Dispatch ID:\n",
            "44e98935-19bb-427e-b96c-019317469fc9\n",
            "Fine tuning new model .......................................................................................................................\n",
            "New Service ID: 66fc231550e1b9bb8cb4cb94\n"
          ]
        }
      ],
      "source": [
        "task = \"Fine-tuning an LLM to translate a sentence without any emojis into a string of only emojis with roughly the same meaning.\"\n",
        "for k in main_agent.submit(\n",
        "    task=task,\n",
        "    name=\"Fine-Tuned Emoji Translator\",\n",
        "    data_format=\"[sentence] | [matching emoji string]\"\n",
        "):\n",
        "    print(k.decode(), end=\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {},
      "outputs": [],
      "source": [
        "emoji_agent = cc.get_deployment(\"66fc231550e1b9bb8cb4cb94\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Swing dancing with my cat | 🐱💃'"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Swing dancing with my cat | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Saved up for a brand new car | 🚗💸'"
            ]
          },
          "execution_count": 31,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Saved up for a brand new car | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Paint me a scenic mountain, Mr. Ross | 🏔️🌲'"
            ]
          },
          "execution_count": 32,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Paint me a scenic mountain, Mr. Ross | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\"Let's grab some burgers on the way home | 🍔👌\""
            ]
          },
          "execution_count": 33,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Let's grab some burgers on the way home | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Starting my vacation tomorrow! | 🏖️🌴'"
            ]
          },
          "execution_count": 34,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Starting my vacation tomorrow! | \")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "covalent-cfs",
      "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.10.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}