llm_inference(Qwen2_v_s_Llama3).ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DbvmBSP-7Br4"
      },
      "source": [
        "Install transformers add upgrade transformers:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CTgUk9-Rcyal"
      },
      "outputs": [],
      "source": [
        "!pip install transformers torch accelerate bitsandbytes sentencepiece gradio\n",
        "!pip install --upgrade transformers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J1v3gE6a7EXy"
      },
      "source": [
        "This is the full import section:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qk1dlL2WnR1F"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig, pipeline\n",
        "from huggingface_hub import login\n",
        "from google.colab import userdata\n",
        "import gradio as gr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gUzF5yVp7H3O"
      },
      "source": [
        "Hugging Face Login, modelid and config, put the secret to HF_TOKEN in your notebook secret, please remeber to request the llama model files access permission."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JV1isDhZ7LRI"
      },
      "outputs": [],
      "source": [
        "#print(userdata.get('HF_TOKEN'))\n",
        "login(token=userdata.get('HF_TOKEN'), add_to_git_credential=True)\n",
        "\n",
        "\n",
        "#model_id = 'meta-llama/Meta-Llama-3.1-8B-Instruct'\n",
        "model_id = 'Qwen/Qwen2-7B-Instruct'\n",
        "\n",
        "#quantization_config\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_quant_type='nf4',\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16\n",
        ")\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "config = AutoConfig.from_pretrained(model_id)\n",
        "config.rope_scaling = { \"type\": \"linear\", \"factor\": 2.0 }  # Adjust the factor as needed\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map='auto')\n",
        "#model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qwPwZYZh7PPv"
      },
      "source": [
        "Text generator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rw2VYBWb7ZIY"
      },
      "outputs": [],
      "source": [
        "text_generator = pipeline(\n",
        "    'text-generation',\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    max_new_tokens=1024,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aGexDSKR7Xko"
      },
      "source": [
        "Get response:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sQNzwWxJ7sLm"
      },
      "outputs": [],
      "source": [
        "def get_response(prompt):\n",
        "  response = text_generator(prompt)\n",
        "  return response[0]['generated_text']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VRHuYIfC8KIq"
      },
      "source": [
        "Give a test:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-JHAwQ18Mid"
      },
      "outputs": [],
      "source": [
        "prompt = \"write a java code, output the difference between current time and 2001/7/1\"\n",
        "response = get_response(prompt)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Add web GUI access"
      ],
      "metadata": {
        "id": "Aqn9vFeChoZY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# use Gradio to create web\n",
        "def gradio_interface(prompt):\n",
        "    return get_response(prompt)\n",
        "\n",
        "\n",
        "interface = gr.Interface(\n",
        "    fn=gradio_interface,\n",
        "    inputs=\"text\",\n",
        "    outputs=\"text\",\n",
        "    title=\"LLM inference 8b in notebook\",\n",
        "    description=\"please input the prompt text.\"\n",
        ")\n",
        "\n",
        "# start it up\n",
        "interface.launch(share=True)"
      ],
      "metadata": {
        "id": "BzLS_bWhiXt_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Clean up the resources if necessary"
      ],
      "metadata": {
        "id": "tVH_pnzphwPX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del model  # remove model\n",
        "del tokenizer # remove tokenizer\n",
        "torch.cuda.empty_cache()\n"
      ],
      "metadata": {
        "id": "GbrccUDahsEH"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Conclusion:**\n",
        "\n",
        "tried two models 'meta-llama/Meta-Llama-3.1-8B-Instruct' and 'Qwen/Qwen2-7B-Instruct'.\n",
        "\n",
        "\n",
        "*   both of them should be quantized in google colab 16G GPU\n",
        "*   Qwen2's performance is better than Llama3.1 in Chinese\n",
        "*   It seems Qwen2 is faster than Llama3.1, 25 seconds v.s. 1.5 minutes. But both of them are slow, maybe the stream output will have a better uer experence.\n",
        "*   Llama3.1 will repeat the output, maybe need to be optimized.\n",
        "*   After quantization, the memory usage around 4G, the GPU usage around 6G.\n",
        "\n"
      ],
      "metadata": {
        "id": "Ptc2imCRmtGP"
      }
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}