generate_sage_github.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyMiyKh6T/fZZEZW1Kahmw4C",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/alga-hopf/alpaca_lora_sage/blob/main/generate_sage_github.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WehxHFXTNar_"
      },
      "outputs": [],
      "source": [
        "!pip3 install transformers peft sentencepiece"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer\n",
        "from peft import PeftModel\n",
        "import torch"
      ],
      "metadata": {
        "id": "ilEDdFA5NdSx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "eeH79-3iNe3z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Some functions"
      ],
      "metadata": {
        "id": "RKDsi_z5dkfY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(prompt_input, input_instr=None, temperature=0.1, top_p=0.75, top_k=40, num_beams=4, max_new_tokens=128):\n",
        "  instruction_format = (\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n ### Instruction:\\n{instruction}\\n\\n### Response:\")\n",
        "  example = {\"instruction\": prompt_input}\n",
        "  prompt  = instruction_format.format_map(example)\n",
        "  inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "  input_ids = inputs[\"input_ids\"].to(device)\n",
        "  generation_config = GenerationConfig(temperature=temperature, top_p=top_p, top_k=top_k, num_beams=num_beams)\n",
        "\n",
        "  with torch.no_grad():\n",
        "      generation_output = model.generate(input_ids=input_ids, generation_config=generation_config, return_dict_in_generate=True, output_scores=True, max_new_tokens=max_new_tokens)\n",
        "\n",
        "  s = generation_output.sequences[0]\n",
        "  output = tokenizer.decode(s)\n",
        "  output = output.split(tokenizer.eos_token)[0][3:]\n",
        "  return output\n",
        "\n",
        "def sage_code_from_output(output):\n",
        "  code = output.split(\"### Response:\")[1]\n",
        "  code = code.splitlines()\n",
        "  sage_code = \"\"\n",
        "  for s in code:\n",
        "    sage_code += s.strip() + \"\\n\"\n",
        "  with open('output_code.sage', 'w') as f:\n",
        "    f.write(sage_code)"
      ],
      "metadata": {
        "id": "BAOg1t9QNiuB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
      ],
      "metadata": {
        "id": "40JHUz4DNnm3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Import the model"
      ],
      "metadata": {
        "id": "E8eA70TWdpXQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint_base =  \"decapoda-research/llama-7b-hf\"\n",
        "out_dir = \"/content/drive/MyDrive/\"\n",
        "checkpoint_lora = out_dir + \"alpaca_lora_sage\"\n",
        "model_max_length = 512\n",
        "tokenizer = LlamaTokenizer.from_pretrained(checkpoint_base, model_max_length=model_max_length, padding_side=\"right\", use_fast=False)\n",
        "DEFAULT_PAD_TOKEN = \"[PAD]\"\n",
        "DEFAULT_EOS_TOKEN = \"</s>\"\n",
        "DEFAULT_BOS_TOKEN = \"<s>\"\n",
        "DEFAULT_UNK_TOKEN = \"<unk>\"\n",
        "tokenizer.pad_token = DEFAULT_PAD_TOKEN\n",
        "tokenizer.eos_token = DEFAULT_EOS_TOKEN\n",
        "tokenizer.bos_token = DEFAULT_BOS_TOKEN\n",
        "tokenizer.unk_token = DEFAULT_UNK_TOKEN\n",
        "model = LlamaForCausalLM.from_pretrained(checkpoint_base, torch_dtype=torch.float32, low_cpu_mem_usage=True, device_map={\"\": device})\n",
        "model = PeftModel.from_pretrained(model, checkpoint_lora, torch_dtype=torch.float32, device_map={\"\": device})"
      ],
      "metadata": {
        "id": "ztcUEWnTNps2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.eval();"
      ],
      "metadata": {
        "id": "VU42335VNr7x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate text\n",
        "Write your prompt in the second line of the next cell and execute."
      ],
      "metadata": {
        "id": "M-Mb5YgMdtkk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "prompt = \"Compute the rank of the matrix A.\"\n",
        "output = evaluate(prompt)"
      ],
      "metadata": {
        "id": "HTFh05e7NvA4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(output)"
      ],
      "metadata": {
        "id": "3TqUHOJzNvDu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The next cell generates a sagemath script containing the output above."
      ],
      "metadata": {
        "id": "NJAQe54G0yA2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sage_code_from_output(output)"
      ],
      "metadata": {
        "id": "DZUvjIUtN0dv"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}