{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TFu_ibC1eYrz"
      },
      "outputs": [],
      "source": [
        "!pip install torch transformers bitsandbytes -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zs7QNs0Tet6r"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import pipeline, BitsAndBytesConfig\n",
        "from IPython.display import clear_output\n",
        "from google.colab import output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lOhAjLdI2oFt"
      },
      "outputs": [],
      "source": [
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_8bit=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v4uIN6uIeyl3"
      },
      "outputs": [],
      "source": [
        "class ChatBot:\n",
        "    _instance = None\n",
        "    _current_model = None\n",
        "\n",
        "    def __init__(self, model_slug=None):\n",
        "        if model_slug and model_slug != ChatBot._current_model:\n",
        "            self.load_model(model_slug)\n",
        "            ChatBot._current_model = model_slug\n",
        "\n",
        "        self.messages = []\n",
        "        self.max_tokens = 2048\n",
        "        self.temperature = 0.5\n",
        "        self.top_k = 100\n",
        "        self.top_p = 0.95\n",
        "\n",
        "    @classmethod\n",
        "    def get_instance(cls, model_slug=None):\n",
        "        if not cls._instance or (model_slug and model_slug != cls._current_model):\n",
        "            cls._instance = cls(model_slug)\n",
        "        return cls._instance\n",
        "\n",
        "    def load_model(self, model_slug):\n",
        "        print(f\"Loading model {model_slug}...\")\n",
        "        self.pipeline = pipeline(\n",
        "            \"text-generation\",\n",
        "            model=model_slug,\n",
        "            model_kwargs={\"quantization_config\": quantization_config},\n",
        "            device_map=\"auto\",\n",
        "        )\n",
        "        clear_output()\n",
        "        print(\"Model loaded successfully!\")\n",
        "\n",
        "    def reset_conversation(self, system_message):\n",
        "        \"\"\"Reset the conversation with a new system message\"\"\"\n",
        "        self.messages = [{\"role\": \"system\", \"content\": system_message}]\n",
        "\n",
        "    def get_response(self, user_input):\n",
        "        \"\"\"Get response with current parameters\"\"\"\n",
        "        self.messages.append({\"role\": \"user\", \"content\": user_input})\n",
        "        outputs = self.pipeline(\n",
        "            self.messages,\n",
        "            max_new_tokens=self.max_tokens,\n",
        "            do_sample=True,\n",
        "            temperature=self.temperature,\n",
        "            top_k=self.top_k,\n",
        "            top_p=self.top_p\n",
        "        )\n",
        "        response = outputs[0][\"generated_text\"][-1]\n",
        "        content = response.get('content', 'No content available')\n",
        "        self.messages.append({\"role\": \"assistant\", \"content\": content})\n",
        "        return content\n",
        "\n",
        "    def update_params(self, max_tokens=None, temperature=None, top_k=None, top_p=None):\n",
        "        \"\"\"Update generation parameters\"\"\"\n",
        "        if max_tokens is not None:\n",
        "            self.max_tokens = max_tokens\n",
        "        if temperature is not None:\n",
        "            self.temperature = temperature\n",
        "        if top_k is not None:\n",
        "            self.top_k = top_k\n",
        "        if top_p is not None:\n",
        "            self.top_p = top_p"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H2n_6Xcue3Vn"
      },
      "outputs": [],
      "source": [
        "def run_chatbot(\n",
        "    model=None,\n",
        "    system_message=\"You are Orca Mini, You are expert in following given instructions, Think step by step before coming up with final answer\",\n",
        "    max_tokens=None,\n",
        "    temperature=None,\n",
        "    top_k=None,\n",
        "    top_p=None,\n",
        "):\n",
        "    try:\n",
        "        # Get or create chatbot instance\n",
        "        chatbot = ChatBot.get_instance(model)\n",
        "\n",
        "        # Update parameters if provided\n",
        "        chatbot.update_params(max_tokens, temperature, top_k, top_p)\n",
        "\n",
        "        # Reset conversation with new system message\n",
        "        chatbot.reset_conversation(system_message)\n",
        "\n",
        "        print(\"Chatbot: Hi! Type 'quit' to exit.\")\n",
        "\n",
        "        while True:\n",
        "            user_input = input(\"You: \").strip()\n",
        "            if user_input.lower() == 'quit':\n",
        "                break\n",
        "            try:\n",
        "                response = chatbot.get_response(user_input)\n",
        "                print(\"Chatbot:\", response)\n",
        "            except Exception as e:\n",
        "                print(f\"Chatbot: An error occurred: {str(e)}\")\n",
        "                print(\"Please try again.\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Error in chatbot: {str(e)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "JEqgoAH2fC6h"
      },
      "outputs": [],
      "source": [
        "run_chatbot(model=\"pankajmathur/orca_mini_phi-4\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tGW8wsfAfHDf"
      },
      "outputs": [],
      "source": [
        "# # change system message\n",
        "# run_chatbot(\n",
        "#     system_message=\"You are Orca Mini, You are expert in logic, Think step by step before coming up with final answer\",\n",
        "#     max_tokens=1024,\n",
        "#     temperature=0.3\n",
        "# )"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}