{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# LogSage Inference in Colab\n",
        "\n",
        "Run this notebook with a GPU runtime.\n",
        "\n",
        "It loads:\n",
        "\n",
        "- base model: `unsloth/Qwen2.5-7B-Instruct-bnb-4bit`\n",
        "- adapter: `auro-rirum/LogSage-Qwen2.5-7B-QLoRA-v0`\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install -U \"bitsandbytes>=0.46.1\" accelerate peft transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from peft import PeftModel\n",
        "from transformers import AutoModelForCausalLM\n",
        "\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\"unsloth/Qwen2.5-7B-Instruct-bnb-4bit\")\n",
        "model = PeftModel.from_pretrained(base_model, \"auro-rirum/LogSage-Qwen2.5-7B-QLoRA-v0\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
        "from peft import PeftModel\n",
        "\n",
        "base_model_id = \"unsloth/Qwen2.5-7B-Instruct-bnb-4bit\"\n",
        "adapter_id = \"auro-rirum/LogSage-Qwen2.5-7B-QLoRA-v0\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)\n",
        "\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        ")\n",
        "\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    base_model_id,\n",
        "    quantization_config=bnb_config,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True,\n",
        ")\n",
        "\n",
        "model = PeftModel.from_pretrained(\n",
        "    base_model,\n",
        "    adapter_id,\n",
        ")\n",
        "\n",
        "model.eval()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "print(torch.cuda.is_available())\n",
        "print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"No GPU\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\n",
        "        \"role\": \"system\",\n",
        "        \"content\": \"You are LogSage, a log-analysis assistant. Return only valid JSON with these keys: issue, root_cause, severity, fix, confidence.\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"\"\"Analyze this log and return only the JSON diagnosis:\n",
        "\n",
        "{\"level\": \"error\", \"component\": \"database\", \"message\": \"Connection attempt failed\", \"details\": \"psycopg2.connection: host=db.internal port=5432 timeout=10s\"}\"\"\"\n",
        "    }\n",
        "]\n",
        "\n",
        "prompt = tokenizer.apply_chat_template(\n",
        "    messages,\n",
        "    tokenize=False,\n",
        "    add_generation_prompt=True\n",
        ")\n",
        "\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
        "\n",
        "outputs = model.generate(\n",
        "    **inputs,\n",
        "    max_new_tokens=250,\n",
        "    do_sample=False,\n",
        "    pad_token_id=tokenizer.eos_token_id\n",
        ")\n",
        "\n",
        "generated = outputs[0][inputs[\"input_ids\"].shape[-1]:]\n",
        "print(tokenizer.decode(generated, skip_special_tokens=True))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}