{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Snapgate-3B Inference Notebook\n",
        "Notebook untuk menjalankan `kadalicious22/snapgate-3B` di Jupyter/Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install -q transformers accelerate torch sentencepiece bitsandbytes\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "MODEL_NAME='kadalicious22/snapgate-3B'\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\n",
        "    MODEL_NAME,\n",
        "    trust_remote_code=True\n",
        ")\n",
        "\n",
        "# Gunakan load_in_4bit=True jika VRAM kecil\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    MODEL_NAME,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map='auto',\n",
        "    trust_remote_code=True\n",
        ")\n",
        "\n",
        "print('Model loaded successfully')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def chat(prompt, max_new_tokens=512, temperature=0.7):\n",
        "    messages = [\n",
        "        {'role':'system','content':'You are Snapgate assistant.'},\n",
        "        {'role':'user','content':prompt}\n",
        "    ]\n",
        "\n",
        "    text = tokenizer.apply_chat_template(\n",
        "        messages,\n",
        "        tokenize=False,\n",
        "        add_generation_prompt=True\n",
        "    )\n",
        "\n",
        "    inputs = tokenizer(\n",
        "        text,\n",
        "        return_tensors='pt'\n",
        "    ).to(model.device)\n",
        "\n",
        "    outputs = model.generate(\n",
        "        **inputs,\n",
        "        max_new_tokens=max_new_tokens,\n",
        "        temperature=temperature,\n",
        "        do_sample=True\n",
        "    )\n",
        "\n",
        "    reply = tokenizer.decode(\n",
        "        outputs[0][inputs.input_ids.shape[-1]:],\n",
        "        skip_special_tokens=True\n",
        "    )\n",
        "    return reply\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Test prompt\n",
        "print(chat('Jelaskan AI agent dalam bahasa Indonesia'))\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}