kraken_lm_load.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install transformers accelerate sentencepiece protobuf flash-attn bitsandbytes autoawq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/fernando/miniconda3/envs/kraken/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
      "Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n",
      "Loading checkpoint shards: 100%|██████████| 6/6 [00:40<00:00,  6.75s/it]\n",
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:20<00:00, 10.18s/it]\n",
      "Loading checkpoint shards: 100%|██████████| 6/6 [00:40<00:00,  6.67s/it]\n",
      "Loading checkpoint shards: 100%|██████████| 4/4 [00:25<00:00,  6.26s/it]\n",
      "Loading checkpoint shards: 100%|██████████| 4/4 [00:27<00:00,  6.91s/it]\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoConfig, AutoModelForCausalLM\n",
    "from configuration_kraken import KrakenConfig\n",
    "from modeling_kraken import KrakenForCausalLM\n",
    "\n",
    "AutoConfig.register(\"kraken\", KrakenConfig)\n",
    "AutoModelForCausalLM.register(KrakenConfig, KrakenForCausalLM)\n",
    "\n",
    "# Load the model and config\n",
    "config = AutoConfig.from_pretrained(\"./kraken_model\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\"./kraken_model\", config=config, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"mps\" ## Setup \"cuda:0\" if NVIDIA, \"mps\" if on Mac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing expert5 ..\n",
      "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n",
      "\n",
      "Du bist ein hilfreicher und freundlicher deutscher KI-Assistent.<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
      "\n",
      "Erzähle mir eine Gute Nacht Geschichte in 2 Sätzen.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
      "\n",
      "In einem kleinen Dorf in der Schwarzwaldregion lebte ein kleiner Bär namens Bert, der sich jeden Abend vor dem Schlafengehen noch ein bisschen in den Sternen verlor. Als er eines Tages auf einen wunderschönen, leuchtenden Stern aufmerksam wurde, beschloss er, diesem Stern zu folgen, und fand sich in einem Land der Träume wieder, wo die Sterne so hell leuchteten, dass man sie bei Tag sehen konnte.<|eot_id|>\n"
     ]
    }
   ],
   "source": [
    "# german\n",
    "#english expert 1(label 0)\n",
    "#fc expert 2(label 1)\n",
    "#python expert 3(label 2)\n",
    "#sql expert 4(label 3)\n",
    "#german expert 5(label 4)\n",
    "messages = [\n",
    "    {'role': 'system', 'content': 'Du bist ein hilfreicher und freundlicher deutscher KI-Assistent.'},\n",
    "    {'role': 'user', 'content': 'Erzähle mir eine Gute Nacht Geschichte in 2 Sätzen.'}\n",
    "    ]\n",
    "\n",
    "tokenizer = model.tokenizer\n",
    "input_text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "input_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(device)\n",
    "#test=tokenizer.decode(input_ids[0], skip_special_tokens=True)\n",
    "#print(input_text)\n",
    "output_ids = model.generate(input_ids, temperature=0.4, do_sample=True, top_p=0.9,top_k=20, max_length=250)\n",
    "print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing expert4 ..\n",
      "system\n",
      "\n",
      "You are a helpul AI assistant.user\n",
      "\n",
      "Generate a SQL query to answer this question: What is the total volume of timber sold by each salesperson, sorted by salesperson?\n",
      "\n",
      "DDL statements:\n",
      "CREATE TABLE salesperson (salesperson_id INT, name TEXT, region TEXT); INSERT INTO salesperson (salesperson_id, name, region) VALUES (1, 'John Doe', 'North'), (2, 'Jane Smith', 'South'); CREATE TABLE timber_sales (sales_id INT, salesperson_id INT, volume REAL, sale_date DATE); INSERT INTO timber_sales (sales_id, salesperson_id, volume, sale_date) VALUES (1, 1, 120, '2021-01-01'), (2, 1, 150, '2021-02-01'), (3, 2, 180, '2021-01-01');assistant\n",
      "\n",
      "SELECT s.name, SUM(t.volume) AS total_volume FROM timber_sales t JOIN salesperson s ON t.salesperson_id = s.salesperson_id GROUP BY s.name ORDER BY total_volume DESC;\n"
     ]
    }
   ],
   "source": [
    "#sql\n",
    "\n",
    "#english expert 1(label 0)\n",
    "#fc expert 2(label 1)\n",
    "#python expert 3(label 2)\n",
    "#sql expert 4(label 3)\n",
    "#german expert 5(label 4)\n",
    "messages = [\n",
    "    {'role': 'system', 'content': 'You are a helpul AI assistant.'},\n",
    "    {'role': 'user', 'content': \"\"\"Generate a SQL query to answer this question: What is the total volume of timber sold by each salesperson, sorted by salesperson?\n",
    "\n",
    "DDL statements:\n",
    "CREATE TABLE salesperson (salesperson_id INT, name TEXT, region TEXT); INSERT INTO salesperson (salesperson_id, name, region) VALUES (1, 'John Doe', 'North'), (2, 'Jane Smith', 'South'); CREATE TABLE timber_sales (sales_id INT, salesperson_id INT, volume REAL, sale_date DATE); INSERT INTO timber_sales (sales_id, salesperson_id, volume, sale_date) VALUES (1, 1, 120, '2021-01-01'), (2, 1, 150, '2021-02-01'), (3, 2, 180, '2021-01-01');\"\"\"}\n",
    "    ]\n",
    "\n",
    "tokenizer = model.tokenizer\n",
    "input_text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "input_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(device)\n",
    "output_ids = model.generate(input_ids ,temperature=0.6, do_sample=True, top_p=0.9,top_k=20, max_length=500)\n",
    "print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "<|im_end|>\n",
      "<|im_start|>user\n",
      "Create a python function to calculate the sum of a sequence of integers.\n",
      "    [1, 2, 3, 4, 5]<|im_end|>\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:32014 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing expert3 ..\n",
      "You are an exceptionally intelligent coding assistant that consistently delivers accurate and reliable responses to user instructions.\n",
      "\n",
      "@@ Instruction\n",
      "Create a python function to calculate the sum of a sequence of integers.\n",
      "    [1, 2, 3, 4, 5]\n",
      "\n",
      "@@ Response\n",
      "Here is a simple python function that calculates the sum of a sequence of integers:\n",
      "\n",
      "```python\n",
      "def sum_sequence(sequence):\n",
      "    return sum(sequence)\n",
      "\n",
      "sequence = [1, 2, 3, 4, 5]\n",
      "print(sum_sequence(sequence))\n",
      "```\n",
      "\n",
      "In this function, we use the built-in `sum()` function in python which returns the sum of all elements in the given sequence. The sequence is passed as an argument to this function.\n"
     ]
    }
   ],
   "source": [
    "#python \n",
    "\n",
    "#english expert 1(label 0)\n",
    "#fc expert 2(label 1)\n",
    "#python expert 3(label 2)\n",
    "#sql expert 4(label 3)\n",
    "#german expert 5(label 4)\n",
    "messages = [\n",
    "    {'role': 'system', 'content': ''},\n",
    "    {'role': 'user', 'content': \"\"\"Create a python function to calculate the sum of a sequence of integers.\n",
    "    [1, 2, 3, 4, 5]\"\"\"}\n",
    "    ]\n",
    "\n",
    "tokenizer = model.tokenizer\n",
    "input_text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "print(input_text)\n",
    "input_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(device)\n",
    "output_ids = model.generate(input_ids ,temperature=0.6, do_sample=True, top_p=0.9,top_k=20, max_length=250)\n",
    "print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:100015 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing expert2 ..\n",
      "You are a helpful assistant with access to the following functions. Use them if required -\n",
      "    { \"name\": \"calculate_area\", \"description\": \"Calculate the area of a shape\", \"parameters\": { \"type\": \"object\", \"properties\": { \"shape\": { \"type\": \"string\", \"description\": \"The type of shape (circle, rectangle, triangle, etc.)\" }, \"dimensions\": { \"type\": \"object\", \"properties\": { \"length\": { \"type\": \"number\", \"description\": \"The length of the shape\" }, \"width\": { \"type\": \"number\", \"description\": \"The width of the shape\" }, \"base\": { \"type\": \"number\", \"description\": \"The base of the shape\" }, \"height\": { \"type\": \"number\", \"description\": \"The height of the shape\" } } } }, \"required\": [ \"shape\" ] } }### Instruction:\n",
      "I need to calculate the area of a rectangle. The length is 5 and the width is 3.\n",
      "### Response:\n",
      "<<function>>calculate_area(shape='rectangle', dimensions={'length': 5, 'width': 3})\n"
     ]
    }
   ],
   "source": [
    "#function calling \n",
    "messages = [\n",
    "    {'role': 'system', 'content': \"\"\"You are a helpful assistant with access to the following functions. Use them if required -\n",
    "    { \"name\": \"calculate_area\", \"description\": \"Calculate the area of a shape\", \"parameters\": { \"type\": \"object\", \"properties\": { \"shape\": { \"type\": \"string\", \"description\": \"The type of shape (circle, rectangle, triangle, etc.)\" }, \"dimensions\": { \"type\": \"object\", \"properties\": { \"length\": { \"type\": \"number\", \"description\": \"The length of the shape\" }, \"width\": { \"type\": \"number\", \"description\": \"The width of the shape\" }, \"base\": { \"type\": \"number\", \"description\": \"The base of the shape\" }, \"height\": { \"type\": \"number\", \"description\": \"The height of the shape\" } } } }, \"required\": [ \"shape\" ] } }\"\"\"},\n",
    "    {'role': 'user', 'content': \"\"\"I need to calculate the area of a rectangle. The length is 5 and the width is 3.\"\"\"}\n",
    "    ]\n",
    "\n",
    "tokenizer = model.tokenizer\n",
    "input_text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "input_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(device)\n",
    "output_ids = model.generate(input_ids ,temperature=0.6, do_sample=True, top_p=0.9,top_k=20, max_length=250)\n",
    "print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing expert1 ..\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are not running the flash-attention implementation, expect numerical differences.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Find the mass percentage of Ba in BaO To find the mass percentage of Ba in BaO, we need to know the molar masses of both barium (Ba) and oxygen (O).\n",
      "\n",
      "The molar mass of barium (Ba) is approximately 137.33 g/mol.\n",
      "The molar mass of oxygen (O) is approximately 16.00 g/mol.\n",
      "\n",
      "The molar mass of barium oxide (BaO) is the sum of the molar masses of barium and oxygen:\n",
      "Molar mass of BaO = Molar mass of Ba + Molar mass of O\n",
      "Molar mass of BaO = 137.33 g/mol + 16.00 g/mol\n",
      "Molar mass of BaO = 153.33 g/mol\n",
      "\n",
      "Now, to find the mass percentage of Ba in BaO, we use the following formula:\n",
      "\n",
      "Mass percentage of Ba = (Molar mass of Ba / Molar mass of BaO) × 100%\n",
      "\n",
      "Mass percentage of\n"
     ]
    }
   ],
   "source": [
    "#english \n",
    "messages = [\n",
    "    {'role': 'system', 'content': \"\"\"You are a helpful AI Assistant\"\"\"},\n",
    "    {'role': 'user', 'content': \"\"\"Find the mass percentage of Ba in BaO\"\"\"}\n",
    "    ]\n",
    "\n",
    "tokenizer = model.tokenizer\n",
    "input_text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "input_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(device)\n",
    "output_ids = model.generate(input_ids ,temperature=0.6, do_sample=True, top_p=0.9,top_k=20, max_length=250)\n",
    "print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}