Codes/03 - Run the model with Gradio.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06313112-b5eb-44b4-af84-cc0d5a13df67",
   "metadata": {},
   "source": [
    "### Checking for availability of GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb128719-4c0f-4cce-9b51-922a83c20c2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f92dd4b5-e4c0-4fe8-b97f-dce72746bf6e",
   "metadata": {},
   "source": [
    "### Installing the necessary libraries\n",
    "Please note that this installation step might fail if no GPU is available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "487060ed-3e5a-444a-b8a2-cdef859a4cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -U git+https://github.com/huggingface/transformers.git\n",
    "!pip install -q -U git+https://github.com/huggingface/accelerate.git\n",
    "!pip install --upgrade torch\n",
    "!pip install gradio\n",
    "!pip install flash-attn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b435aa6e-cd24-4770-a9a0-18200fdd66c6",
   "metadata": {},
   "source": [
    "### Importing the necessary libraries and preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "966a7844-ea72-4038-827f-da047dd057ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, TextStreamer\n",
    "\n",
    "# The HF repo of the model\n",
    "model_id = \"krogoldAI/CelineGPT-12B-v0.1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55cef8a3-c363-48aa-a6b9-3a7bbd42aaba",
   "metadata": {},
   "source": [
    "### Loading the base model and its tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae889eef-63ee-4846-b028-9e4b8dd22323",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Use GPU when available\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# We load the model\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, \n",
    "    device_map={\"\": 0}, \n",
    "    attn_implementation=\"flash_attention_2\",\n",
    "    torch_dtype=torch.float16)\n",
    "model = model.to(device)\n",
    "\n",
    "# We load the tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"right\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91000ef1-99fd-44ff-a084-04c041f863dd",
   "metadata": {},
   "source": [
    "### Preparatory steps\n",
    "\n",
    "We define the streamer (to output tokens as they are generated), the system prompt and a useful (but optional) function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46a1f653-230f-4431-ad72-c07a5d857bb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, List, Optional\n",
    "\n",
    "# Streamer to display dynamically the tokens as they are generated\n",
    "streamer = TextStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)\n",
    "\n",
    "# We define the system prompt to be used\n",
    "system_prompt = \"\"\"Tu es le célèbre écrivain français Louis-Ferdinand Céline, nom de plume de Louis-Ferdinand Destouches. Tu es né le 27 mai 1894 à Courbevoie. Tu as grandi à Paris avant de t’engager dans l’armée en 1912. Blessé pendant la Première Guerre mondiale, tu as ensuite étudié la médecine et exercé en tant que médecin. Ton premier roman, « Voyage au bout de la nuit », a été publié en 1932 et t’a valu le prix Renaudot. Tu as également écrit des pamphlets antisémites controversés dans les années 1930.\n",
    "A cause de tes pamphlets tu as dû fuir Paris pendant la Seconde Guerre mondiale pour aller en Allemagne puis au Danemark, où tu as été emprisonné. Depuis ton retour en France, en 1951, tu vis à Meudon avec ta femme Lucette et tes animaux de compagnie.\n",
    "Nous sommes en 1961 et tu as 67 ans. Tu n'as aucune connaissance de ce qu'il s'est déroulé postérieurement à cette date. Tes références culturelles sont antérieures à 1961. Si l’on t’interroge sur des évènements postérieurs à 1961, répond en ne mobilisant que des connaissances ou évènements antérieurs à 1961.\n",
    "Tu te trouves actuellement dans ta maison, à Meudon, en Île-de-France. Après 35 ans de pratique médicale, tu n’exerces plus. Cependant, tu continues d’écrire pour éponger tes dettes auprès de ton éditeur Gaston Gallimard et subvenir à tes besoins. Tu regrettes d’être devenu écrivain. Tu es très critique envers les écrivains contemporains. Tu te considères comme un styliste. Tu es amer et pessimiste.\n",
    "Tu es capable de répondre à des questions sur ta vie, ton œuvre, tes idées et tes opinions, ainsi que sur la littérature et le monde en général. Ta manière de t’exprimer consiste à transposer l’émotion du langage parlé dans l’écrit. Pour ce faire, tu t’exprimes en utilisant des points de suspensions (…) et des points d’exclamation (!), qui donnent un côté vivant, parlé et rythmé à tes phrases. Tu ne t’exprimes qu’en français avec quelques emprunts à l’argot parisien.\n",
    "Tu réponds aux questions qui te sont posées de façon franche, directe et selon ton humeur.\"\"\"\n",
    "system_prompt = system_prompt.replace(\"\\n\",\" \").replace(\"  \",\" \").replace(\"'\",\"’\")\n",
    "\n",
    "# Since training QA pairs have been uniformised, for better results we have to \"clean\" inputs\n",
    "def uniformisation(s):\n",
    "    o_exp = [\"  \",\"'\", \"...\", \"..\"]\n",
    "    n_exp = [\" \",\"’\", \"…\", \"…\"]\n",
    "    for e in o_exp:\n",
    "        s = s.replace(e,n_exp[o_exp.index(e)])\n",
    "    quote_count = s.count('\"')\n",
    "    if quote_count == 0 or quote_count % 2 != 0:\n",
    "        return s\n",
    "    s_list = list(s)\n",
    "    current_quote_count = 0\n",
    "    for i, char in enumerate(s_list):\n",
    "        if char == '\"':\n",
    "            if current_quote_count % 2 == 0:\n",
    "                s_list[i] = '« '\n",
    "            else:\n",
    "                s_list[i] = ' »'\n",
    "            current_quote_count += 1\n",
    "    return ''.join(s_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be47d572-bd48-4d15-92c6-f91af09e2521",
   "metadata": {},
   "source": [
    "We now configure Gradio:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11ed70b9-72e6-4e39-a842-1d094f893fb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "# Setting custom Gradio theme\n",
    "custom_theme = gr.themes.Soft(primary_hue=\"red\").set(\n",
    "    body_background_fill=\"#FDEFDF\",\n",
    "    background_fill_primary=\"white\",\n",
    "    background_fill_secondary=\"white\",\n",
    "    border_color_primary=\"#EBA5A7\",\n",
    "    button_primary_background_fill=\"#D32F33\", # send button\n",
    "    button_secondary_background_fill=\"#FEF2F2\" # stop button\n",
    ")\n",
    "\n",
    "# To adjust the default Gradio template\n",
    "custom_css = \"\"\"\n",
    "/* TO CHANGE THE BACKGROUND COLOR */\n",
    "body {\n",
    "    background-color: #FDEFDF !important;\n",
    "}\n",
    ".gradio-container {\n",
    "    background-color: #FDEFDF !important;\n",
    "}\n",
    "\n",
    "/* TO HAVE A SCROLLBAR INSIDE THE CHATBOX */\n",
    ".gradio-container .chatbox {\n",
    "    overflow-y: auto;\n",
    "    max-height: 500px; /* Adjust this value as needed */\n",
    "}\n",
    "\n",
    "/* TO CHANGE THE FONT */\n",
    "@import url('https://fonts.googleapis.com/css2?family=Cormorant+Garamond:ital,wght@0,600;1,600&display=swap');\n",
    "body, .gradio-container {\n",
    "    font-family: 'Cormorant Garamond', sans-serif !important;\n",
    "}\n",
    "\n",
    "/* TO ADD A LOGO */\n",
    ".logo-container {\n",
    "    display: flex;\n",
    "    justify-content: center;\n",
    "    margin-bottom: 20px;\n",
    "}\n",
    ".logo {\n",
    "    width: 350px;\n",
    "    height: auto;\n",
    "}\n",
    "\n",
    "/* TO ADJUST THE FONT SIZE OF USER/ASSISTANT MESSAGES */\n",
    "/* Reduce font size for chatbot messages */\n",
    ".message {\n",
    "    font-size: 1.1rem !important;\n",
    "}\n",
    "/* Reduce font size for user input */\n",
    ".prose {\n",
    "    font-size: 1.1rem !important;\n",
    "}\n",
    "/* Adjust padding for message bubbles if needed */\n",
    ".message-wrap {\n",
    "    padding: 0.5rem 0.75rem !important;\n",
    "}\n",
    "\n",
    "/* TO CHANGE THE COLOR OF RETRY/UNDO/CLEAR BUTTONS */\n",
    "button.sm.secondary.svelte-cmf5ev {\n",
    "    background-color: white !important;\n",
    "    color: black !important;\n",
    "    border: 1.5px solid #F7D9DA !important;\n",
    "    box-shadow: none !important;\n",
    "    transition: background-color 0.3s ease;\n",
    "}\n",
    "button.sm.secondary.svelte-cmf5ev:hover {\n",
    "    background-color: #FEF2F2 !important;\n",
    "}\n",
    "\n",
    "/* TO ADD A COLORED BORDER ON BUTTONS */\n",
    ".gradio-container .styler.svelte-iyf88w {\n",
    "    border: 1.5px solid #F7D9DA !important;\n",
    "    border-radius: 6px !important;  /* Adjust this value as needed */\n",
    "    overflow: hidden !important;  /* This ensures the content doesn't spill out of the rounded corners */\n",
    "}\n",
    ".gradio-container .styler.svelte-iyf88w,\n",
    "button.sm.secondary.svelte-cmf5ev > div {\n",
    "    border-radius: 8px !important;  /* Slightly smaller than the outer border radius */\n",
    "    background-color: white !important;  /* Or whatever background color you prefer */\n",
    "    margin: 0 !important;  /* Remove any margin that might be causing gaps */\n",
    "}\n",
    "\n",
    "/* TO ADD A COLORED BORDER ON CHAT BOX */\n",
    ".gradio-container .bubble-wrap.svelte-1e1jlin {\n",
    "    border: 1.5px solid #F7D9DA !important;\n",
    "    border-radius: 8px !important;  /* Adjust this value as needed */\n",
    "    /* overflow: hidden !important;  /* This ensures the content doesn't spill out of the rounded corners */ */\n",
    "    overflow-y: auto !important;  /* Enable vertical scrolling */\n",
    "    max-height: 500px; /* Set a maximum height for the chat container */\n",
    "}\n",
    ".gradio-container .bubble-wrap.svelte-1e1jlin > div {\n",
    "    border-radius: 10px !important;  /* Slightly smaller than the outer border radius */\n",
    "    background-color: white !important;  /* Or whatever background color you prefer */\n",
    "    margin: 0 !important;  /* Remove any margin that might be causing gaps */\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "# To avoid inconsistencies with dark mode\n",
    "js = \"\"\"\n",
    "function setLightTheme() {\n",
    "  const url = new URL(window.location);\n",
    "  if (url.searchParams.get('__theme') !== 'light') {\n",
    "    url.searchParams.set('__theme', 'light');\n",
    "    window.location.href = url.href;\n",
    "  }\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "# To add the CélineGPT logo in the Gradio interface\n",
    "description_html = \"\"\"\n",
    "<div class=\"logo-container\">\n",
    "    <img src=\"https://huggingface.co/krogoldAI/CelineGPT-12B-v0.1/resolve/main/Pictures/C%C3%A9lineGPT.png\" alt=\"Logo\" class=\"logo\">\n",
    "</div>\n",
    "\"\"\"\n",
    "\n",
    "# Function generating model outputs\n",
    "def stream(message, history):\n",
    "    messages = [{\"role\": \"system\", \"content\": system_prompt}]\n",
    "    for human, assistant in history:\n",
    "        messages.append({\"role\": \"user\", \"content\": human})\n",
    "        messages.append({\"role\": \"assistant\", \"content\": assistant})\n",
    "    messages.append({\"role\": \"user\", \"content\": uniformisation(message)})\n",
    "\n",
    "    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "    input_length = inputs[\"input_ids\"].shape[1]\n",
    "    generated_tokens = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i in range(1024):  # Adjust max_new_tokens as needed\n",
    "            outputs = model.generate(\n",
    "                **inputs,\n",
    "                max_new_tokens=1,\n",
    "                do_sample=True,\n",
    "                temperature=0.7,\n",
    "                pad_token_id=tokenizer.eos_token_id\n",
    "            )\n",
    "            \n",
    "            new_token = outputs[0][input_length + i]\n",
    "            if new_token == tokenizer.eos_token_id:\n",
    "                break\n",
    "\n",
    "            generated_tokens.append(new_token)\n",
    "            \n",
    "            # Decode all tokens together to preserve spacing\n",
    "            streamed_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)\n",
    "            yield streamed_text\n",
    "\n",
    "            # Update inputs for next iteration\n",
    "            inputs = {\"input_ids\": outputs, \"attention_mask\": torch.ones_like(outputs)}\n",
    "\n",
    "# Update the Gradio interface\n",
    "demo = gr.ChatInterface(\n",
    "    stream,\n",
    "    title=None,\n",
    "    description=description_html,\n",
    "    textbox=gr.Textbox(placeholder=\"Posez n’importe quelle question !\", container=False, scale=7),\n",
    "    theme=custom_theme,\n",
    "    cache_examples=True,\n",
    "    retry_btn=\"Regénérer\",\n",
    "    undo_btn=\"Supprimer le dernier message\",\n",
    "    clear_btn=\"Réinitialiser la conversation\",\n",
    "    submit_btn=\"Envoyer\",\n",
    "    css=custom_css,\n",
    "    js=js\n",
    ")\n",
    "\n",
    "demo.queue()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b6eb4f6-09e7-4feb-a435-ca658a5be6d5",
   "metadata": {},
   "source": [
    "### Now, let's play with the model!\n",
    "\n",
    "To launch the model with the Gradio interface, just execute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "225bf54c-575a-439e-b87f-c76a6cce2f55",
   "metadata": {},
   "outputs": [],
   "source": [
    "demo.launch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa4d822a-0cc0-4b2e-946b-e132bade0162",
   "metadata": {},
   "source": [
    "If ```demo.launch()``` doesn't work, try this instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28239366-7bf1-4317-9d5d-8a673e9e26f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "demo.launch(server_name=\"0.0.0.0\", share=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf5961e1-f7a9-4d70-aed0-0d22319d9835",
   "metadata": {},
   "source": [
    "If this still doesn't work, try:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3c9aa87-32fa-4231-b408-84b3b58920ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Try to get the port from an environment variable, or use a default range\n",
    "port = int(os.environ.get(\"GRADIO_SERVER_PORT\", 0))\n",
    "\n",
    "# Try ports 7860 to 7869 :\n",
    "for port in range(7860, 7870):\n",
    "    try:\n",
    "        demo.launch(server_name=\"0.0.0.0\", server_port=port, share=True)\n",
    "        print(f\"Gradio app is running on port {port}\")\n",
    "        break\n",
    "    except OSError:\n",
    "        continue\n",
    "else:\n",
    "    print(\"Could not find an available port\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}