{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checking for availability of GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "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,
   "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 flash-attn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing the necessary libraries and preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "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",
   "metadata": {},
   "source": [
    "### Loading the base model and its tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    ").to(device)\n",
    "\n",
    "# We load the tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"right\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting generation configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "executionInfo": {
     "elapsed": 421,
     "status": "ok",
     "timestamp": 1716824737228,
     "user": {
      "displayName": "Hintergrundstrahlung",
      "userId": "11564469484860642689"
     },
     "user_tz": -120
    },
    "id": "mFKHmBa8tfiB"
   },
   "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)\n",
    "\n",
    "# Function to handle multi-turn chat mode with history of conversation\n",
    "def chat(\n",
    "    query: str,\n",
    "    history: Optional[List[Dict]] = None,\n",
    "    temperature: float = 0.8,\n",
    "    top_p: float = 1.0,\n",
    "    top_k: float = 0,\n",
    "    repetition_penalty: float = 1.2,\n",
    "    max_new_tokens: int = 1024,\n",
    "    **kwargs,\n",
    "):\n",
    "    if history is None:\n",
    "        history = [{\"role\": \"system\", \"content\": system_prompt}]\n",
    "\n",
    "    query = uniformisation(query)\n",
    "    history.append({\"role\": \"user\", \"content\": query})\n",
    "\n",
    "    input_ids = tokenizer.apply_chat_template(history, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\n",
    "    input_length = input_ids.shape[1]\n",
    "\n",
    "    generated_outputs = model.generate(\n",
    "        input_ids=input_ids, \n",
    "        generation_config=GenerationConfig(\n",
    "            temperature=temperature,\n",
    "            do_sample=temperature > 0.0, # i.e. do_sample = True\n",
    "            top_p=top_p,\n",
    "            top_k=top_k,\n",
    "            repetition_penalty=repetition_penalty,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            pad_token_id=tokenizer.pad_token_id,\n",
    "            **kwargs,\n",
    "        ),\n",
    "        streamer=streamer,\n",
    "        return_dict_in_generate=True,\n",
    "        num_return_sequences=1,\n",
    "        pad_token_id=tokenizer.pad_token_id\n",
    "    )\n",
    "\n",
    "    generated_tokens = generated_outputs.sequences[0, input_length:]\n",
    "    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)\n",
    "\n",
    "    history.append({\"role\": \"assistant\", \"content\": generated_text})\n",
    "\n",
    "    return generated_text, history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's play with the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "historique = None\n",
    "while True:\n",
    "    user_input = input(\"Moi :\")\n",
    "    if user_input.lower() == \"exit\":\n",
    "        break\n",
    "    print(\"L.-F. Céline :\", end=\"\")\n",
    "    reponse, historique = chat(user_input, historique)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "historique"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyPxeWDit8rdz3qUjoGzVwiO",
   "gpuType": "V100",
   "provenance": [
    {
     "file_id": "16gchCLhnxhGVLLTVAwII1kXMlcEgxjsc",
     "timestamp": 1713448473332
    },
    {
     "file_id": "1CT3CsG4yKGoP-E2MqUozPwN_1QGwKgLZ",
     "timestamp": 1712924241559
    }
   ]
  },
  "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"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "20b2263fc2e643bf9658d9e1fafd77a6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "44fec00182714f16ab5bc72e6a9f3c5f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7b2a4c7216f043fa814b2610d0ce9b2d",
      "placeholder": "​",
      "style": "IPY_MODEL_20b2263fc2e643bf9658d9e1fafd77a6",
      "value": "Loading checkpoint shards: 100%"
     }
    },
    "45ed85d5f08c4dd6bd17fc7eefa8598d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_44fec00182714f16ab5bc72e6a9f3c5f",
       "IPY_MODEL_963e7589963d46d3b9db0455f901078e",
       "IPY_MODEL_c651060afcfa4836a6a75408b15183c3"
      ],
      "layout": "IPY_MODEL_8a8bd07d3f9942998032f32763b3a863"
     }
    },
    "70551b8d0a9b491fa7fab953e4e2ee4c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "7b2a4c7216f043fa814b2610d0ce9b2d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8a8bd07d3f9942998032f32763b3a863": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "91395c8d0bb54db482af1fc18d24e796": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "963e7589963d46d3b9db0455f901078e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_91395c8d0bb54db482af1fc18d24e796",
      "max": 3,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_70551b8d0a9b491fa7fab953e4e2ee4c",
      "value": 3
     }
    },
    "a8d255d330634d72b031fba1576b22e3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "c651060afcfa4836a6a75408b15183c3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_fc71e182277a4c879205b6bbfab75837",
      "placeholder": "​",
      "style": "IPY_MODEL_a8d255d330634d72b031fba1576b22e3",
      "value": " 3/3 [00:05<00:00,  1.75s/it]"
     }
    },
    "fc71e182277a4c879205b6bbfab75837": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}