zero-data-model-foundry.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Covalent Cloud\n",
        "\n",
        "☁️ [Covalent Cloud](https://www.covalent.xyz/cloud/)\n",
        "\n",
        "🚀 [Covalent Cloud QuickStart](https://docs.covalent.xyz/docs/cloud/cloud_quickstart)\n",
        "\n",
        "### Covalent OS\n",
        "\n",
        "🌟 [GitHub: Covalent Open-Source](https://github.com/AgnostiqHQ/covalent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Building A Zero-Data Model Foundry\n",
        "\n",
        "## Setting Up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "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": 57,
      "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": 25,
      "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": 26,
      "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": 27,
      "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=400px/>\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": 28,
      "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=\"7 hours\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "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",
        "    \"\"\"Initialize the service that host the data generator LLM\"\"\"\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",
        "    \"\"\"Endpoint to generate data using the LLM\"\"\"\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=400px/>\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": 30,
      "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=\"7 hours\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {},
      "outputs": [],
      "source": [
        "@cc.service(executor=ft_service_ex, volume=volume, name=\"Custom Fine-Tuned Model\")\n",
        "def finetuned_llm_service(ft_model_path):\n",
        "\n",
        "    \"\"\"Serves a newly 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_length=100):\n",
        "\n",
        "    \"\"\"Generate text from a prompt using the fine-tuned language model.\"\"\"\n",
        "\n",
        "    output = pipe(prompt, truncation=True, max_length=max_length, num_return_sequences=1)\n",
        "    return output[0][\"generated_text\"]\n",
        "\n",
        "\n",
        "@finetuned_llm_service.endpoint(\"/stream\", streaming=True)\n",
        "def generate_stream(model, tokenizer, prompt, prepend_prompt=False, max_tokens=100):\n",
        "\n",
        "    \"\"\"Prompt Llama-like model to stream generated text.\"\"\"\n",
        "\n",
        "    def _starts_with_space(_tokenizer, _token_id):\n",
        "        token = _tokenizer.convert_ids_to_tokens(_token_id)\n",
        "        return token.startswith('▁')\n",
        "\n",
        "    _input = tokenizer(prompt, return_tensors='pt')\n",
        "    _input = _input.to(\"cuda\")\n",
        "\n",
        "    if prepend_prompt:\n",
        "        yield prompt\n",
        "\n",
        "    for output_length in range(max_tokens):\n",
        "        output = model.generate(**_input, max_new_tokens=1)\n",
        "        current_token_id = output[0][-1]\n",
        "        if current_token_id == tokenizer.eos_token_id:\n",
        "            break\n",
        "\n",
        "        current_token = tokenizer.decode(\n",
        "            current_token_id, skip_special_tokens=True\n",
        "        )\n",
        "        if _starts_with_space(tokenizer, current_token_id.item()) and output_length > 1:\n",
        "            current_token = ' ' + current_token\n",
        "        yield current_token\n",
        "\n",
        "        _input = {\n",
        "            'input_ids': output.to(\"cuda\"),\n",
        "            'attention_mask': torch.ones(1, len(output[0])).to(\"cuda\"),\n",
        "        }"
      ]
    },
    {
      "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=400px/>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [],
      "source": [
        "agent_ex = cc.CloudExecutor(env=FT_ENV, num_cpus=12, memory=\"12GB\", time_limit=\"7 hours\")\n",
        "\n",
        "@cc.service(executor=agent_ex, name=\"Fine Tuner Agent\", auth=False, volume=volume)\n",
        "def agent(lattice, llm_api):\n",
        "\n",
        "    \"\"\"Initialize the agent. Not much to do, just store the input params.\"\"\"\n",
        "\n",
        "    return {\"finetune_lattice\": lattice, \"llm_api\": llm_api}\n",
        "\n",
        "\n",
        "@agent.endpoint(\"/submit\", streaming=True)\n",
        "def submit(\n",
        "    finetune_lattice, llm_api,\n",
        "    *,\n",
        "    task,\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 = llm_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",
        "    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=400px/>\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": 33,
      "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",
        "- H100 GPU\n",
        "- 32 GB RAM\n",
        "\n",
        "Tasks exit and release resources after completion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "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.H100,\n",
        "    memory=\"32GB\",\n",
        "    time_limit=\"7 hours\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "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": 36,
      "metadata": {},
      "outputs": [],
      "source": [
        "cpu_ex = cc.CloudExecutor(env=FT_ENV, num_cpus=12, memory=\"12GB\", time_limit=\"4 hours\")\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": 37,
      "metadata": {},
      "outputs": [],
      "source": [
        "@ct.lattice(executor=cpu_ex, workflow_executor=cpu_ex)\n",
        "def setup_workflow(ft_workflow, data_generator_model=\"unsloth/llama-3-8b-Instruct\"):\n",
        "\n",
        "    \"\"\"Set up for everything.\"\"\"\n",
        "\n",
        "    data_generator_handle = llm_data_generator(data_generator_model)\n",
        "    agent_handle = agent(ft_workflow, data_generator_handle)\n",
        "    return agent_handle"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "446da31768a749d1a05f3d936449263c",
              "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:  79405a20-2cc9-453c-9ae0-f8f4dd4b049e\n",
            "Main Agent Service ID:  66460db7f7d37dbf2a4689cb\n"
          ]
        }
      ],
      "source": [
        "dispatch_id = cc.dispatch(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": 39,
      "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 Tuner Agent                                                    │\n",
            "│  Description   Initialize the agent. Not much to do, just store the input params.  │\n",
            "│  Function ID   66460db7f7d37dbf2a4689cb                                            │\n",
            "│  Address       https://fn.prod.covalent.xyz/066460db7f7d37dbf2a4689cb              │\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(\"66460db7f7d37dbf2a4689cb\")\n",
        "print(main_agent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Example 1: Spoiler Detector"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating fine-tuning data .......................................................................\n",
            "Generated 3014 total examples\n",
            "Saving dataset at /volumes/model-storage/data_3014-a804ac4f-3b7d-4b61-8857-6e90ec6419a8\n",
            "\n",
            "Dispatching fine-tuning workflow\n",
            "Dispatch ID:\n",
            "bb8c0ae6-6584-499b-89c1-e862ddb69239\n",
            "Fine tuning new model ....................................................................\n",
            "New Service ID: 664610f3f7d37dbf2a4689cf\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(task=task, min_new_examples=3000):\n",
        "    print(k.decode(), end=\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Load a client for the new spoiler agent\n",
        "spoiler_agent = cc.get_deployment(\"664610f3f7d37dbf2a4689cf\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "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. ## NO SPOILER'"
            ]
          },
          "execution_count": 59,
          "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": 60,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'SPOILER'"
            ]
          },
          "execution_count": 60,
          "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": 61,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'SPOILER'"
            ]
          },
          "execution_count": 61,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spoiler_agent.generate(prompt=\"Soylent green is people! ##\").split(\"##\")[-1].strip()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'SPOILER'"
            ]
          },
          "execution_count": 62,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spoiler_agent.generate(prompt=\"Anyways, at the end you find out that the Planet of The Apes was Earth all along. ##\").split(\"##\")[-1].strip()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Example 2: Grammar Corrector"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating fine-tuning data ........................................\n",
            "Generated 2000 total examples\n",
            "Saving dataset at /volumes/model-storage/data_2000-a3e2a820-bd58-4dff-aa40-8c3954ccab7a\n",
            "\n",
            "Dispatching fine-tuning workflow\n",
            "Dispatch ID:\n",
            "27000d76-343d-4674-9f27-c554e7cea481\n",
            "Fine tuning new model .............................................................\n",
            "New Service ID: 6646144af7d37dbf2a4689d3\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(task=task, data_format=data_format):\n",
        "    print(k.decode(), end=\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {},
      "outputs": [],
      "source": [
        "grammar_agent = cc.get_deployment(\"6646144af7d37dbf2a4689d3\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "I should have never got a pet.\n",
            "Jerry and I argued about it.\n",
            "He said, 'No cat bites its own tail.'\n",
            "But 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.\")  # dogs, their, tails"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Example 3: Emoji translation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating fine-tuning data .................................................\n",
            "Generated 2000 total examples\n",
            "Saving dataset at /volumes/model-storage/data_2000-e089dd31-cb29-4f42-baff-1e23ef4b579f\n",
            "\n",
            "Dispatching fine-tuning workflow\n",
            "Dispatch ID:\n",
            "81c30729-c84c-493a-a1a2-c2915aab5cdd\n",
            "Fine tuning new model ..............................................................\n",
            "New Service ID: 66461753f7d37dbf2a4689d7\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(task=task, data_format=\"[sentence] | [matching emoji string]\"):\n",
        "    print(k.decode(), end=\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {},
      "outputs": [],
      "source": [
        "emoji_agent = cc.get_deployment(\"66461753f7d37dbf2a4689d7\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Dancing with my cat | 🐈💕'"
            ]
          },
          "execution_count": 66,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Dancing with my cat | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Got a brand new car | 🚗💸'"
            ]
          },
          "execution_count": 67,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Got a brand new car | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Paint me a scenic mountain, Mr. Ross | 🏔️'"
            ]
          },
          "execution_count": 68,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Paint me a scenic mountain, Mr. Ross | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "\"Let's eat! | 🍴👅\""
            ]
          },
          "execution_count": 69,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Let's eat! | \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Burning Man | 🔥'"
            ]
          },
          "execution_count": 70,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emoji_agent.generate(prompt=\"Burning Man | \")"
      ]
    }
  ],
  "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.14"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}