Sarvam-M-Test.ipynb

sarvam-m-test.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyO+I2SFYQc2PYVx8LqyvGjj",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/hathibelagal-dev/1932446cb69ca14ea1d908721fbdd3eb/sarvam-m-test.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Sarvam-M is an LLM that's signficantly better at working with Indian languages.\n",
        "\n",
        "Here's a notebook you can use to chat with it.\n",
        "\n",
        "-- [Hathibelagal](https://x.com/hathibel)"
      ],
      "metadata": {
        "id": "h7Xf41HF0ePO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip3 install bitsandbytes"
      ],
      "metadata": {
        "id": "ZuxZPxFYzdLb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n9PpSSZEzQSa"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
        "from transformers import TextStreamer\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = \"sarvamai/sarvam-m\""
      ],
      "metadata": {
        "id": "AyUKtt7OzaQr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "q_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
        "    bnb_4bit_quant_type=\"fp4\"\n",
        ")"
      ],
      "metadata": {
        "id": "q0evG_pn2bD-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_name, torch_dtype=torch.bfloat16, device_map=\"cuda\",\n",
        "    quantization_config=q_config\n",
        ")"
      ],
      "metadata": {
        "id": "rN0lSvK8zT6c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "tQ-mnkP07COQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"कितने वेद हैं?\"\n",
        "\n",
        "messages = [\n",
        "    {\n",
        "        \"role\": \"system\",\n",
        "        \"content\": \"You're a very helpful Indian assistant.\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": prompt\n",
        "    },\n",
        "]\n",
        "\n",
        "model_input = tokenizer.apply_chat_template(\n",
        "    messages, tokenize=True, add_generation_prompt=True,\n",
        "    return_tensors=\"pt\").to(\"cuda\")"
      ],
      "metadata": {
        "id": "gW4HQw7u1DYB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "streamer = TextStreamer(tokenizer)"
      ],
      "metadata": {
        "id": "ZWji14DR7WCe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "outputs = model.generate(model_input, max_new_tokens=1024, streamer=streamer)"
      ],
      "metadata": {
        "id": "dGbiTWAX5WtT"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}