API tools faq
paste
Advertisement
Guest User

Untitled

a guest
Apr 9th, 2023
209
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 40.82 KB | None | 0 0
raw download clone embed print report
  1. {
  2. "nbformat": 4,
  3. "nbformat_minor": 0,
  4. "metadata": {
  5. "colab": {
  6. "provenance": []
  7. },
  8. "kernelspec": {
  9. "name": "python3",
  10. "display_name": "Python 3"
  11. },
  12. "language_info": {
  13. "name": "python"
  14. },
  15. "gpuClass": "standard",
  16. "accelerator": "GPU"
  17. },
  18. "cells": [
  19. {
  20. "cell_type": "markdown",
  21. "source": [
  22. "Specify model and group size, as well as whether to save the model to drive (by default, `/drive/My Drive/neko_models/`) when finished and load the model from drive on subsequent runs. If you don't have enough drive space to download the model or you're running this somewhere without drive compatibility, leave it off.\n",
  23. "\n",
  24. "A https://huggingface.co/ URL to a quantized model must be provided (path to the actual checkpoint (so the URL to the safetensors, pt, or bin in the repository must be provided. It also must be compatible with GPTQ-for-Llama Triton branch, i.e., not CUDA, otherwise it won't work)."
  25. ],
  26. "metadata": {
  27. "id": "hxMbk3dZiy3x"
  28. }
  29. },
  30. {
  31. "cell_type": "code",
  32. "execution_count": null,
  33. "metadata": {
  34. "id": "blPQFwhrRsFE"
  35. },
  36. "outputs": [],
  37. "source": [
  38. "use_drive = True #@param {type:\"boolean\"}\n",
  39. "\n",
  40. "num_params = \"13b\" #@param [\"7b\", \"13b\", \"30b\", \"65b\"]\n",
  41. "\n",
  42. "group_size = \"128\" #@param [\"32\", \"128\"]\n",
  43. "group_size = int(group_size)\n",
  44. "\n",
  45. "huggingface_url = '' #@param {type:\"string\"}\n",
  46. "\n",
  47. "group_and_model = '/'.join(huggingface_url[len('https://huggingface.co/'):].split('/')[:2])"
  48. ]
  49. },
  50. {
  51. "cell_type": "markdown",
  52. "source": [
  53. "Download repository, install necessary Python packages, retrieve model checkpoint."
  54. ],
  55. "metadata": {
  56. "id": "0w7eC5-mkoN6"
  57. }
  58. },
  59. {
  60. "cell_type": "code",
  61. "source": [
  62. "#if use_runpod:\n",
  63. "# !apt-get -y update\n",
  64. "# !apt-get -y install python3.10-dev\n",
  65. "# !python -m pip install --upgrade pip\n",
  66. "\n",
  67. "!git clone https://github.com/qwopqwop200/GPTQ-for-LLaMa.git\n",
  68. "%cd 'GPTQ-for-LLaMa'\n",
  69. "!git checkout 29c0e4a4552c26a8a18abff305504d2cc86be539\n",
  70. "!pip install -r requirements.txt\n",
  71. "%cd ..\n",
  72. "\n",
  73. "import os\n",
  74. "import sys\n",
  75. "\n",
  76. "weights_name = os.path.basename(huggingface_url)\n",
  77. "drive_weights_path = os.path.join('drive/My Drive/neko_models/{}'.format(weights_name))\n",
  78. "\n",
  79. "if use_drive:\n",
  80. " # Connect to drive.\n",
  81. " from google.colab import drive\n",
  82. " drive.mount('/content/drive')\n",
  83. " weights_path = drive_weights_path\n",
  84. "else:\n",
  85. " weights_path = weights_name\n",
  86. "\n",
  87. "if not use_drive or not os.path.exists(drive_weights_path):\n",
  88. " # Download model, save to drive if flag is set.\n",
  89. " !wget {huggingface_url}\n",
  90. " if use_drive:\n",
  91. " if not os.path.exists('drive/My Drive/neko_models/'):\n",
  92. " os.mkdir('drive/My Drive/neko_models/')\n",
  93. " !cp {weights_name} \"{drive_weights_path}\"\n",
  94. "\n",
  95. "sys.path.insert(0, 'GPTQ-for-LLaMa/')"
  96. ],
  97. "metadata": {
  98. "id": "MYKNE67ciiqc"
  99. },
  100. "execution_count": null,
  101. "outputs": []
  102. },
  103. {
  104. "cell_type": "markdown",
  105. "source": [
  106. "Load model."
  107. ],
  108. "metadata": {
  109. "id": "h-yldfa3eFTR"
  110. }
  111. },
  112. {
  113. "cell_type": "code",
  114. "source": [
  115. "import time\n",
  116. "\n",
  117. "import torch\n",
  118. "import torch.nn as nn\n",
  119. "\n",
  120. "from gptq import *\n",
  121. "from modelutils import *\n",
  122. "from quant import *\n",
  123. "\n",
  124. "from transformers import AutoTokenizer\n",
  125. "\n",
  126. "DEV = torch.device('cuda')\n",
  127. "\n",
  128. "def load_quant(model, checkpoint, wbits, groupsize, device):\n",
  129. " from transformers import LlamaConfig, LlamaForCausalLM \n",
  130. " config = LlamaConfig.from_pretrained(model)\n",
  131. " def noop(*args, **kwargs):\n",
  132. " pass\n",
  133. " torch.nn.init.kaiming_uniform_ = noop \n",
  134. " torch.nn.init.uniform_ = noop \n",
  135. " torch.nn.init.normal_ = noop \n",
  136. "\n",
  137. " torch.set_default_dtype(torch.half)\n",
  138. " transformers.modeling_utils._init_weights = False\n",
  139. " torch.set_default_dtype(torch.half)\n",
  140. " model = LlamaForCausalLM(config)\n",
  141. " torch.set_default_dtype(torch.float)\n",
  142. " model = model.eval()\n",
  143. " layers = find_layers(model)\n",
  144. " for name in ['lm_head']:\n",
  145. " if name in layers:\n",
  146. " del layers[name]\n",
  147. " make_quant(model, layers, wbits, groupsize)\n",
  148. "\n",
  149. " print('Loading model ...')\n",
  150. " if checkpoint.endswith('.safetensors'):\n",
  151. " from safetensors import safe_open\n",
  152. " with safe_open(checkpoint, framework=\"pt\", device=0) as f:\n",
  153. " for k in f.keys():\n",
  154. " layer = {}\n",
  155. " layer[k] = f.get_tensor(k)\n",
  156. " model.load_state_dict(layer, strict=False)\n",
  157. " del(layer)\n",
  158. " else:\n",
  159. " model.load_state_dict(torch.load(checkpoint))\n",
  160. " print('Done.')\n",
  161. "\n",
  162. " return model\n",
  163. "\n",
  164. "model = load_quant(group_and_model, weights_path, 4, group_size, DEV)\n",
  165. "model.to(DEV)\n",
  166. "tokenizer = AutoTokenizer.from_pretrained(group_and_model, use_fast=False)"
  167. ],
  168. "metadata": {
  169. "id": "v7NFb4Q5fR-J"
  170. },
  171. "execution_count": null,
  172. "outputs": []
  173. },
  174. {
  175. "cell_type": "markdown",
  176. "source": [
  177. "Define generation methods."
  178. ],
  179. "metadata": {
  180. "id": "Q_G_lmDseJat"
  181. }
  182. },
  183. {
  184. "cell_type": "code",
  185. "source": [
  186. "import torch\n",
  187. "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
  188. "\n",
  189. "def stm_next_tokens(model, tokenizer, tokenized, context_len, max_gen_len,\n",
  190. " mask_id, temperature=0.8, top_p=0.95, tfs=1.0, typical=1.0,\n",
  191. " penalty_range=1024, penalty_slope=0.7, penalty=1.1,\n",
  192. " past_key_values=None, return_past_key_values=True):\n",
  193. " #tokenized = tokenizer.encode(inp, return_tensors='pt').to(DEV)\n",
  194. " total_len = min(context_len, max_gen_len + tokenized.shape[1])\n",
  195. "\n",
  196. " tokens = torch.full((1, total_len), mask_id).to(DEV)\n",
  197. " tokens[0, :tokenized.shape[1]] = tokenized[0]\n",
  198. "\n",
  199. " output_past_key_values = None\n",
  200. "\n",
  201. " if past_key_values and return_past_key_values:\n",
  202. " output_past_key_values = past_key_values\n",
  203. "\n",
  204. " for cur_id in range(tokenized.shape[1], total_len):\n",
  205. " #print(cur_id - tokenized.shape[1])\n",
  206. " if past_key_values:\n",
  207. " output = model(tokens[:, cur_id-1:cur_id], past_key_values=past_key_values, use_cache=True)\n",
  208. " else:\n",
  209. " output = model(tokens[:, :cur_id], use_cache=True)\n",
  210. "\n",
  211. " if not past_key_values:\n",
  212. " logits = output.logits[:, cur_id-1, :]\n",
  213. " if return_past_key_values:\n",
  214. " output_past_key_values = output.past_key_values\n",
  215. " else:\n",
  216. " logits = output.logits[:, 0, :]\n",
  217. " \n",
  218. " past_key_values = output.past_key_values\n",
  219. " input_ids = tokens[:, cur_id-1:cur_id]\n",
  220. "\n",
  221. " # Apply samplers - do greedy sampling if temperature is 0.\n",
  222. " if temperature > 0:\n",
  223. " next_token_scores = sample_top_p_actual(input_ids, logits,\n",
  224. " top_p)\n",
  225. " next_token_scores = sample_tail_free(input_ids,\n",
  226. " next_token_scores, tfs)\n",
  227. " next_token_scores = sample_typical(input_ids, next_token_scores,\n",
  228. " typical)\n",
  229. " next_token_scores = sample_temperature(input_ids,\n",
  230. " next_token_scores,\n",
  231. " temperature)\n",
  232. " next_token_scores = sample_advanced_repetition_penalty(input_ids,\n",
  233. " next_token_scores,\n",
  234. " penalty_range,\n",
  235. " penalty_slope,\n",
  236. " penalty)\n",
  237. "\n",
  238. " next_token_scores = torch.nn.functional.softmax(next_token_scores,\n",
  239. " dim=-1)\n",
  240. "\n",
  241. " next_token = torch.multinomial(next_token_scores,\n",
  242. " num_samples=1).squeeze(1)\n",
  243. " else:\n",
  244. " next_token = torch.argmax(logits, axis=-1)[0]\n",
  245. "\n",
  246. " tokens[0, cur_id] = next_token\n",
  247. " yield next_token, None\n",
  248. "\n",
  249. " if next_token.item() == tokenizer.eos_token_id:\n",
  250. " yield None, output_past_key_values\n",
  251. " return\n",
  252. " \n",
  253. " yield None, output_past_key_values\n",
  254. " return\n",
  255. "\n",
  256. "# taken from Kobold and transformers so this stuff is AGPL I guess\n",
  257. "def sample_temperature(input_ids, scores, tempt):\n",
  258. " scores = scores / tempt\n",
  259. " return scores\n",
  260. "\n",
  261. "def sample_typical(input_ids, scores, typical, filter_value = -float(\"Inf\"),\n",
  262. " min_tokens_to_keep = 1):\n",
  263. " if filter_value >= 1.0:\n",
  264. " return scores\n",
  265. "\n",
  266. " probs = scores.softmax(dim=-1)\n",
  267. " log_probs = probs.log()\n",
  268. "\n",
  269. " neg_entropy = (probs * log_probs).nansum(dim=-1, keepdim=True)\n",
  270. "\n",
  271. " entropy_deviation = (neg_entropy - log_probs).abs()\n",
  272. "\n",
  273. " _, sorted_indices = torch.sort(entropy_deviation)\n",
  274. " sorted_logits = probs.gather(-1, sorted_indices)\n",
  275. " sorted_indices_to_remove = sorted_logits.cumsum(dim=-1) >= typical\n",
  276. " sorted_indices_to_remove = sorted_indices_to_remove.roll(1, dims=-1)\n",
  277. "\n",
  278. " min_tokens_to_keep = max(min_tokens_to_keep, 1)\n",
  279. " # Keep at least min_tokens_to_keep\n",
  280. " sorted_indices_to_remove[..., : min_tokens_to_keep] = 0\n",
  281. "\n",
  282. " indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n",
  283. " scores = scores.masked_fill(indices_to_remove, filter_value)\n",
  284. " return scores \n",
  285. "\n",
  286. "def sample_top_p_actual(input_ids, scores, top_p, filter_value = -float(\"Inf\"),\n",
  287. " min_tokens_to_keep = 1):\n",
  288. " sorted_logits, sorted_indices = torch.sort(scores, descending=False)\n",
  289. " cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)\n",
  290. "\n",
  291. " # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)\n",
  292. " sorted_indices_to_remove = cumulative_probs <= (1 - top_p)\n",
  293. " if min_tokens_to_keep > 1:\n",
  294. " # Keep at least min_tokens_to_keep\n",
  295. " sorted_indices_to_remove[..., -min_tokens_to_keep :] = 0\n",
  296. "\n",
  297. " # scatter sorted tensors to original indexing\n",
  298. " indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices,\n",
  299. " sorted_indices_to_remove)\n",
  300. " scores = scores.masked_fill(indices_to_remove, filter_value)\n",
  301. " return scores\n",
  302. "\n",
  303. "def sample_advanced_repetition_penalty(input_ids, scores, penalty_range,\n",
  304. " penalty_slope, penalty):\n",
  305. " penalty_range = int(penalty_range)\n",
  306. " clipped_penalty_range = min(input_ids.shape[-1], penalty_range)\n",
  307. "\n",
  308. " if penalty != 1.0:\n",
  309. " if penalty_range > 0:\n",
  310. " if clipped_penalty_range < input_ids.shape[1]:\n",
  311. " input_ids = input_ids[..., -clipped_penalty_range:]\n",
  312. "\n",
  313. " if penalty_slope != 0:\n",
  314. " _penalty = (torch.arange(penalty_range, dtype=scores.dtype,\n",
  315. " device=scores.device)/(penalty_range - 1)) * 2. - 1\n",
  316. " _penalty = (penalty_slope * _penalty) / (1 + torch.abs(_penalty) * (penalty_slope - 1))\n",
  317. " _penalty = 1 + ((_penalty + 1) / 2).unsqueeze(0) * (penalty - 1)\n",
  318. " penalty = _penalty[..., -clipped_penalty_range:]\n",
  319. "\n",
  320. " score = torch.gather(scores, 1, input_ids)\n",
  321. " score = torch.where(score <= 0, score * penalty, score / penalty)\n",
  322. " scores.scatter_(1, input_ids, score)\n",
  323. "\n",
  324. " return scores \n",
  325. "\n",
  326. "def sample_top_a(input_ids, scores, top_a, filter_value = -float(\"Inf\"),\n",
  327. " min_tokens_to_keep = 1):\n",
  328. " if filter_value >= 1.0:\n",
  329. " return scores\n",
  330. "\n",
  331. " sorted_logits, sorted_indices = torch.sort(scores, descending=True)\n",
  332. " probs = sorted_logits.softmax(dim=-1)\n",
  333. "\n",
  334. " # Remove tokens with probability less than top_a*(max(probs))^2 (token with 0 are kept)\n",
  335. " probs_max = probs[..., 0, None]\n",
  336. " sorted_indices_to_remove = probs < probs_max * probs_max * top_a\n",
  337. "\n",
  338. " if min_tokens_to_keep > 1:\n",
  339. " # Keep at least min_tokens_to_keep\n",
  340. " sorted_indices_to_remove[..., : min_tokens_to_keep] = 0\n",
  341. "\n",
  342. " indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices,\n",
  343. " sorted_indices_to_remove)\n",
  344. " scores = scores.masked_fill(indices_to_remove, filter_value)\n",
  345. " return scores \n",
  346. "\n",
  347. "def sample_tail_free(input_ids, scores, tfs, filter_value = -float(\"Inf\"),\n",
  348. " min_tokens_to_keep = 1):\n",
  349. " if filter_value >= 1.0:\n",
  350. " return scores\n",
  351. " sorted_logits, sorted_indices = torch.sort(scores, descending=True)\n",
  352. " probs = sorted_logits.softmax(dim=-1)\n",
  353. "\n",
  354. " # Compute second derivative normalized CDF\n",
  355. " d2 = probs.diff().diff().abs()\n",
  356. " normalized_d2 = d2 / d2.sum(dim=-1, keepdim=True)\n",
  357. " normalized_d2_cdf = normalized_d2.cumsum(dim=-1)\n",
  358. "\n",
  359. " # Remove tokens with CDF value above the threshold (token with 0 are kept)\n",
  360. " sorted_indices_to_remove = normalized_d2_cdf > tfs\n",
  361. "\n",
  362. " # Centre the distribution around the cutoff as in the original implementation of the algorithm\n",
  363. " sorted_indices_to_remove = torch.cat(\n",
  364. " (\n",
  365. " torch.zeros(scores.shape[0], 1, dtype=torch.bool,\n",
  366. " device=scores.device),\n",
  367. " sorted_indices_to_remove,\n",
  368. " torch.ones(scores.shape[0], 1, dtype=torch.bool,\n",
  369. " device=scores.device),\n",
  370. " ),\n",
  371. " dim=-1,\n",
  372. " )\n",
  373. "\n",
  374. " if min_tokens_to_keep > 1:\n",
  375. " # Keep at least min_tokens_to_keep\n",
  376. " sorted_indices_to_remove[..., : min_tokens_to_keep] = 0\n",
  377. "\n",
  378. " indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices,\n",
  379. " sorted_indices_to_remove)\n",
  380. " scores = scores.masked_fill(indices_to_remove, filter_value)\n",
  381. " return scores"
  382. ],
  383. "metadata": {
  384. "id": "1mmAK2zhi6Ak"
  385. },
  386. "execution_count": null,
  387. "outputs": []
  388. },
  389. {
  390. "cell_type": "markdown",
  391. "source": [
  392. "Main GUI. Includes sampler parameters, memory, and lorebook.\n",
  393. "\n",
  394. "Lorebook tip: multiple keys can be associated with an entry by separating them via commas."
  395. ],
  396. "metadata": {
  397. "id": "KZcyrZdreNGX"
  398. }
  399. },
  400. {
  401. "cell_type": "code",
  402. "source": [
  403. "import ipywidgets as widgets\n",
  404. "from IPython.display import display\n",
  405. "import time\n",
  406. "from enum import Enum\n",
  407. "\n",
  408. "context_size = 2048 #@param {type:\"number\"}\n",
  409. "max_gen_len = 160 #@param {type:\"number\"}\n",
  410. "temperature = 1.0 #@param {type:\"number\"}\n",
  411. "top_p = 0.95 #@param {type:\"number\"}\n",
  412. "tfs = 1.0 #@param {type:\"number\"}\n",
  413. "typical = 1.0 #@param {type:\"number\"}\n",
  414. "penalty_range = 2048 #@param {type:\"number\"}\n",
  415. "penalty_slope = 0.7 #@param {type:\"number\"}\n",
  416. "penalty = 1.1 #@param {type:\"number\"}\n",
  417. "lorebook_token_search_range = 1000 #@param {type:\"number\"}\n",
  418. "kv_cache_size = 0 #@param {type:\"number\"}\n",
  419. "\n",
  420. "input_text_area = widgets.Textarea(placeholder='Enter a prompt...',\n",
  421. " layout=widgets.Layout(width='900px',\n",
  422. " height='600px'))\n",
  423. "model.seqlen = context_size\n",
  424. "send_button = widgets.Button(description='Send')\n",
  425. "undo_button = widgets.Button(description='Undo')\n",
  426. "redo_button = widgets.Button(description='Redo')\n",
  427. "retry_button = widgets.Button(description='Retry')\n",
  428. "prev_retry_button = widgets.Button(description='Previous Retry')\n",
  429. "memory_button = widgets.ToggleButton(description='Memory')\n",
  430. "context_button = widgets.ToggleButton(description='Context')\n",
  431. "\n",
  432. "side_panel = widgets.VBox([send_button, undo_button, redo_button,\n",
  433. " retry_button, prev_retry_button,\n",
  434. " memory_button, context_button])\n",
  435. "\n",
  436. "lorebook_dropdown = widgets.Dropdown(options=['(Empty)'],\n",
  437. " layout=widgets.Layout(width='200px'))\n",
  438. "lorebook_add_button = widgets.Button(description='Add New')\n",
  439. "lorebook_remove_button = widgets.Button(description='Remove')\n",
  440. "lorebook_apply_changes_button = widgets.Button(description='Apply Changes')\n",
  441. "lorebook_key_field = widgets.Text(description='Key')\n",
  442. "lorebook_value_field = widgets.Textarea(description='Value',\n",
  443. " layout=widgets.Layout(height='100px'))\n",
  444. "\n",
  445. "lorebook_panel = widgets.VBox([widgets.HBox([lorebook_dropdown,\n",
  446. " lorebook_key_field,\n",
  447. " lorebook_value_field]),\n",
  448. " widgets.HBox([lorebook_add_button,\n",
  449. " lorebook_remove_button,\n",
  450. " lorebook_apply_changes_button])])\n",
  451. "\n",
  452. "lorebook_accordion = widgets.Accordion(children=[lorebook_panel],\n",
  453. " titles=('Lorebook',),\n",
  454. " selected_index=None)\n",
  455. "lorebook_accordion.set_title(0, 'Lorebook')\n",
  456. "\n",
  457. "main_panel = widgets.HBox([widgets.VBox([input_text_area, lorebook_accordion]),\n",
  458. " side_panel])\n",
  459. "output = widgets.Output()\n",
  460. "\n",
  461. "Mode = Enum('Mode', ['INPUT', 'MEMORY', 'GENERATING', 'CONTEXT'])\n",
  462. "\n",
  463. "class State:\n",
  464. " def __init__(self, pos, mode):\n",
  465. " self.pos = pos\n",
  466. " self.mode = mode\n",
  467. " self.mem = ''\n",
  468. " self.saved_input = ''\n",
  469. " self.kv_dict = {}\n",
  470. " self.kv_queue = []\n",
  471. " self.lorebook_dict = {}\n",
  472. " self.lorebook_entry_index = 0\n",
  473. " self.lorebook_prev_key = ''\n",
  474. " self.lorebook_prev_value = ''\n",
  475. " \n",
  476. " def get_kv(self, id):\n",
  477. " if id in self.kv_dict:\n",
  478. " return self.kv_dict[id]\n",
  479. " else:\n",
  480. " return None\n",
  481. " \n",
  482. " def delete_kv(self, id):\n",
  483. " if id in self.kv_dict:\n",
  484. " self.kv_queue.remove(id)\n",
  485. " del self.kv_dict[id]\n",
  486. "\n",
  487. " def add_kv(self, id, kv):\n",
  488. " if id in self.kv_dict:\n",
  489. " self.kv_queue.remove(id)\n",
  490. " self.kv_queue.insert(0, id)\n",
  491. " elif self.get_len_kv() < kv_cache_size:\n",
  492. " self.kv_queue.insert(0, id)\n",
  493. " self.kv_dict[id] = kv\n",
  494. " \n",
  495. " def remove_last_kv(self):\n",
  496. " if self.kv_queue:\n",
  497. " old_id = self.kv_queue.pop()\n",
  498. " del self.kv_dict[old_id]\n",
  499. " \n",
  500. " def get_len_kv(self):\n",
  501. " return len(self.kv_queue)\n",
  502. " \n",
  503. " def clear_kv(self):\n",
  504. " for id in self.kv_queue:\n",
  505. " del self.kv_dict[id]\n",
  506. " self.kv_queue = []\n",
  507. "\n",
  508. "class Position:\n",
  509. " cur_id = 0\n",
  510. " def __init__(self):\n",
  511. " self.id = Position.cur_id\n",
  512. " self.pred = None\n",
  513. " self.succs = []\n",
  514. " self.succ_idx = -1\n",
  515. " self.text = ''\n",
  516. " Position.cur_id += 1\n",
  517. "\n",
  518. "init_pos = Position()\n",
  519. "cur_state = State(init_pos, Mode.INPUT)\n",
  520. "\n",
  521. "def build_context():\n",
  522. " # When creating the context, first, place the full memory followed by a\n",
  523. " # newline.\n",
  524. " #\n",
  525. " # Next, taking the last (max_seq_len-1-max_gen_len-len(mem)) tokens,\n",
  526. " # place these tokens in the context.\n",
  527. " \n",
  528. " inp_tokenized = tokenizer.encode(input_text_area.value, return_tensors='pt', add_special_tokens=False)[0].tolist()\n",
  529. "\n",
  530. " # Search for keys in lorebook, append results to memory.\n",
  531. "\n",
  532. " lorebook_search_text = tokenizer.decode(inp_tokenized[-lorebook_token_search_range:])\n",
  533. " lorebook_values = []\n",
  534. " for key_str, value in cur_state.lorebook_dict.items():\n",
  535. " if not key_str:\n",
  536. " continue\n",
  537. " \n",
  538. " keys = key_str.split(',')\n",
  539. "\n",
  540. " for key in keys:\n",
  541. " if key.strip() in lorebook_search_text:\n",
  542. " lorebook_values.append(value)\n",
  543. " break\n",
  544. "\n",
  545. " memory = '\\n'.join([cur_state.mem, *lorebook_values])\n",
  546. "\n",
  547. " if memory:\n",
  548. " mem_tokenized = tokenizer.encode(memory + '\\n', return_tensors='pt', add_special_tokens=False)[0].tolist()\n",
  549. " else:\n",
  550. " mem_tokenized = []\n",
  551. "\n",
  552. " num_inp_tokens = max(model.seqlen-1-max_gen_len-len(mem_tokenized), 0)\n",
  553. "\n",
  554. " if num_inp_tokens > 0:\n",
  555. " tokenized = mem_tokenized + inp_tokenized[-num_inp_tokens:]\n",
  556. " elif len(mem_tokenized) > 0:\n",
  557. " num_mem_tokens = model.seqlen-1-max_gen_len\n",
  558. " tokenized = mem_tokenized[-num_mem_tokens:]\n",
  559. " else:\n",
  560. " tokenized = []\n",
  561. "\n",
  562. " tokenized.insert(0, tokenizer.bos_token_id)\n",
  563. " detokenized = tokenizer.decode(tokenized)\n",
  564. " return detokenized\n",
  565. "\n",
  566. "def generate():\n",
  567. " # Create the context and send it to the model, update the text area.\n",
  568. " \n",
  569. " gen_context = build_context()\n",
  570. " retokenized = tokenizer.encode(gen_context, return_tensors='pt').to(DEV)\n",
  571. " prev_num_tokens = len(retokenized[0])\n",
  572. "\n",
  573. " output = ''\n",
  574. " past_key_values = None\n",
  575. " num_characters = 0\n",
  576. "\n",
  577. " with torch.no_grad():\n",
  578. " out_tokens = retokenized[0].tolist()\n",
  579. " gen = stm_next_tokens(model, tokenizer, retokenized, model.seqlen,\n",
  580. " max_gen_len, 1, temperature=temperature, top_p=top_p, tfs=tfs,\n",
  581. " typical=typical, penalty_range=penalty_range,\n",
  582. " penalty_slope=penalty_slope, penalty=penalty,\n",
  583. " past_key_values=cur_state.get_kv(cur_state.pos.id),\n",
  584. " return_past_key_values=kv_cache_size != 0)\n",
  585. " for tkn, pkv in gen:\n",
  586. " if tkn is None:\n",
  587. " past_key_values = pkv\n",
  588. " else:\n",
  589. " out_tokens.append(tkn.item())\n",
  590. " output = tokenizer.decode(out_tokens, skip_special_tokens=True)\n",
  591. " num_characters = len(output) - len(gen_context)\n",
  592. " input_text_area.value = cur_state.pos.text + output[-num_characters:]\n",
  593. " torch.cuda.empty_cache()\n",
  594. " return output[-num_characters:], past_key_values\n",
  595. "\n",
  596. "def on_update_input_text_area(change):\n",
  597. " # Input mode: Destroy all successors in the node list.\n",
  598. " #\n",
  599. " # Memory mode: n/a.\n",
  600. " #\n",
  601. " # Action allowed criterion: state.mode == 'input' or state.mode == 'memory'.\n",
  602. "\n",
  603. " if cur_state.mode == Mode.INPUT and (cur_state.pos.succs or cur_state.get_kv(cur_state.pos.id)) and cur_state.pos.text != input_text_area.value:\n",
  604. " if cur_state.pos.succs:\n",
  605. " del cur_state.pos.succs\n",
  606. " cur_state.pos.succs = []\n",
  607. " cur_state.pos.succ_idx = -1\n",
  608. " update_buttons_visible()\n",
  609. " if cur_state.get_kv(cur_state.pos.id):\n",
  610. " cur_state.delete_kv(cur_state.pos.id)\n",
  611. "\n",
  612. "def send():\n",
  613. " cur_state.pos.text = input_text_area.value\n",
  614. " cur_state.mode = Mode.GENERATING\n",
  615. " update_buttons_visible()\n",
  616. "\n",
  617. " if cur_state.get_len_kv() == kv_cache_size and not cur_state.get_kv(cur_state.pos.id):\n",
  618. " cur_state.remove_last_kv()\n",
  619. "\n",
  620. " generation, past_key_values = generate()\n",
  621. "\n",
  622. " new_succ = Position()\n",
  623. " new_succ.pred = cur_state.pos\n",
  624. " #new_succ.text = input_text_area.value + generation\n",
  625. " new_succ.text = input_text_area.value\n",
  626. " cur_state.pos.succs.append(new_succ)\n",
  627. " cur_state.pos.succ_idx = len(cur_state.pos.succs) - 1\n",
  628. " if past_key_values is not None:\n",
  629. " cur_state.add_kv(cur_state.pos.id, past_key_values)\n",
  630. " \n",
  631. " jump_to(new_succ)\n",
  632. "\n",
  633. " cur_state.mode = Mode.INPUT\n",
  634. " update_buttons_visible()\n",
  635. "\n",
  636. "def send_button_clicked(b):\n",
  637. " # Set text in current node to whatever is in the input area, generate text\n",
  638. " # (setting mode to 'generating' in the meantime), create a new successor at\n",
  639. " # head of list with text, set successor position to it, jump to it.\n",
  640. " #\n",
  641. " # Action allowed criterion: state.mode == 'input'.\n",
  642. "\n",
  643. " send()\n",
  644. " \n",
  645. "def undo_button_clicked(b):\n",
  646. " # Jump to predecessor.\n",
  647. " #\n",
  648. " # Action allowed criterion: state.mode == 'input', state.predecessor != nil.\n",
  649. "\n",
  650. " jump_to(cur_state.pos.pred)\n",
  651. "\n",
  652. "def redo_button_clicked(b):\n",
  653. " # Jump to current successor.\n",
  654. " #\n",
  655. " # Action allowed criterion: state.mode == 'input', state.successor_list !=\n",
  656. " # nil.\n",
  657. "\n",
  658. " jump_to(cur_state.pos.succs[cur_state.pos.succ_idx])\n",
  659. "\n",
  660. "def retry_button_clicked(b):\n",
  661. " # Jump to predecessor, then set successor position to next in the list if\n",
  662. " # it exists and jump to it, otherwise send_button_clicked().\n",
  663. " #\n",
  664. " # Action allowed criterion: state.mode == 'input', state.predecessor != nil.\n",
  665. "\n",
  666. " jump_to(cur_state.pos.pred)\n",
  667. "\n",
  668. " if cur_state.pos.succ_idx < len(cur_state.pos.succs) - 1:\n",
  669. " cur_state.pos.succ_idx += 1\n",
  670. " jump_to(cur_state.pos.succs[cur_state.pos.succ_idx])\n",
  671. " else:\n",
  672. " send()\n",
  673. "\n",
  674. "def prev_retry_button_clicked(b):\n",
  675. " # Jump to predecessor, then set successor position to prev in the list and\n",
  676. " # jump to it.\n",
  677. " #\n",
  678. " # Action allowed criterion: state.mode == 'input', state.predecessor != nil,\n",
  679. " # state.predecessor.succ_idx > 0.\n",
  680. "\n",
  681. " jump_to(cur_state.pos.pred)\n",
  682. " cur_state.pos.succ_idx -= 1\n",
  683. " jump_to(cur_state.pos.succs[cur_state.pos.succ_idx])\n",
  684. "\n",
  685. "def memory_button_clicked(b):\n",
  686. " # Input mode: switch modes to 'memory', save current state.\n",
  687. " #\n",
  688. " # Memory mode: switch modes to 'input', save memory, restore current state.\n",
  689. " #\n",
  690. " # Action allowed criterion: state.mode == 'input' or state.mode == 'memory'.\n",
  691. "\n",
  692. " if cur_state.mode == Mode.INPUT:\n",
  693. " cur_state.mode = Mode.MEMORY\n",
  694. " cur_state.saved_input = input_text_area.value\n",
  695. " input_text_area.value = cur_state.mem\n",
  696. " update_buttons_visible()\n",
  697. " elif cur_state.mode == Mode.MEMORY:\n",
  698. " if cur_state.mem != input_text_area.value:\n",
  699. " cur_state.clear_kv()\n",
  700. " cur_state.mode = Mode.INPUT\n",
  701. " cur_state.mem = input_text_area.value\n",
  702. " input_text_area.value = cur_state.saved_input\n",
  703. " update_buttons_visible()\n",
  704. "\n",
  705. "def context_button_clicked(b):\n",
  706. " # Input mode: switch modes to 'context', save current state.\n",
  707. " #\n",
  708. " # Context mode: switch mode to 'input', restore current state.\n",
  709. " #\n",
  710. " # Action allowed criterion: state.mode == 'input' or state.mode == 'context'.\n",
  711. "\n",
  712. " if cur_state.mode == Mode.INPUT:\n",
  713. " cur_state.mode = Mode.CONTEXT\n",
  714. " cur_state.saved_input = input_text_area.value\n",
  715. " input_text_area.value = build_context()\n",
  716. " update_buttons_visible()\n",
  717. " elif cur_state.mode == Mode.CONTEXT:\n",
  718. " cur_state.mode = Mode.INPUT\n",
  719. " input_text_area.value = cur_state.saved_input\n",
  720. " update_buttons_visible()\n",
  721. "\n",
  722. "def jump_to(pos):\n",
  723. " cur_state.pos = pos\n",
  724. " input_text_area.value = pos.text\n",
  725. " update_buttons_visible()\n",
  726. "\n",
  727. "def update_buttons_visible():\n",
  728. " send_button.disabled = cur_state.mode != Mode.INPUT\n",
  729. " undo_button.disabled = cur_state.mode != Mode.INPUT or not cur_state.pos.pred\n",
  730. " redo_button.disabled = cur_state.mode != Mode.INPUT or not cur_state.pos.succs\n",
  731. " retry_button.disabled = cur_state.mode != Mode.INPUT or not cur_state.pos.pred\n",
  732. " prev_retry_button.disabled = cur_state.mode != Mode.INPUT or not cur_state.pos.pred or not cur_state.pos.pred.succ_idx > 0\n",
  733. " memory_button.disabled = cur_state.mode != Mode.INPUT and cur_state.mode != Mode.MEMORY\n",
  734. " context_button.disabled = cur_state.mode != Mode.INPUT and cur_state.mode != Mode.CONTEXT\n",
  735. " input_text_area.disabled = cur_state.mode == Mode.GENERATING or cur_state.mode == Mode.CONTEXT\n",
  736. " lorebook_key_field.disabled = cur_state.mode == Mode.GENERATING or len(cur_state.lorebook_dict) == 0\n",
  737. " lorebook_value_field.disabled = cur_state.mode == Mode.GENERATING or len(cur_state.lorebook_dict) == 0\n",
  738. " lorebook_add_button.disabled = cur_state.mode == Mode.GENERATING\n",
  739. " lorebook_remove_button.disabled = cur_state.mode == Mode.GENERATING or len(cur_state.lorebook_dict) == 0\n",
  740. " lorebook_apply_changes_button.disabled = cur_state.mode == Mode.GENERATING or len(cur_state.lorebook_dict) == 0 or (lorebook_key_field.value == cur_state.lorebook_prev_key and lorebook_value_field.value == cur_state.lorebook_prev_value)\n",
  741. " lorebook_dropdown.disabled = cur_state.mode == Mode.GENERATING\n",
  742. "\n",
  743. "def apply_changes_to_entry(set_prev_key_to_cur):\n",
  744. " lorebook_cur_key = lorebook_key_field.value\n",
  745. " lorebook_cur_value = lorebook_value_field.value\n",
  746. " if not cur_state.lorebook_prev_key or (cur_state.lorebook_prev_key == lorebook_cur_key and cur_state.lorebook_prev_value == lorebook_cur_value) or cur_state.lorebook_prev_key not in cur_state.lorebook_dict:\n",
  747. " return\n",
  748. "\n",
  749. " if cur_state.lorebook_prev_key != lorebook_cur_key:\n",
  750. " del cur_state.lorebook_dict[cur_state.lorebook_prev_key]\n",
  751. " cur_state.lorebook_dict[lorebook_cur_key] = lorebook_cur_value\n",
  752. " new_options = list(lorebook_dropdown.options)\n",
  753. " lorebook_key_idx = new_options.index(cur_state.lorebook_prev_key)\n",
  754. " new_options[lorebook_key_idx] = lorebook_cur_key\n",
  755. " lorebook_dropdown.options = tuple(new_options)\n",
  756. " lorebook_dropdown.value = lorebook_cur_key\n",
  757. " elif cur_state.lorebook_prev_value != lorebook_cur_value:\n",
  758. " cur_state.lorebook_dict[lorebook_cur_key] = lorebook_cur_value\n",
  759. " \n",
  760. " if set_prev_key_to_cur:\n",
  761. " cur_state.lorebook_prev_key = lorebook_cur_key\n",
  762. " cur_state.lorebook_prev_value = lorebook_cur_value\n",
  763. "\n",
  764. "def add_new_entry():\n",
  765. " new_key = 'Key{}'.format(cur_state.lorebook_entry_index)\n",
  766. " new_val = 'Value{}'.format(cur_state.lorebook_entry_index)\n",
  767. " cur_state.lorebook_dict[new_key] = new_val\n",
  768. " cur_state.lorebook_entry_index += 1\n",
  769. "\n",
  770. " if len(cur_state.lorebook_dict) == 1:\n",
  771. " lorebook_dropdown.options = (new_key,)\n",
  772. " else:\n",
  773. " new_options = list(lorebook_dropdown.options)\n",
  774. " new_options.append(new_key)\n",
  775. " lorebook_dropdown.options = tuple(new_options)\n",
  776. " lorebook_dropdown.value = new_key\n",
  777. "\n",
  778. "def remove_entry():\n",
  779. " removed_key = cur_state.lorebook_prev_key\n",
  780. "\n",
  781. " del cur_state.lorebook_dict[cur_state.lorebook_prev_key]\n",
  782. " cur_state.lorebook_prev_key = ''\n",
  783. " cur_state.lorebook_prev_val = ''\n",
  784. "\n",
  785. " if len(cur_state.lorebook_dict) == 0:\n",
  786. " lorebook_dropdown.options = ('(Empty)',)\n",
  787. " else:\n",
  788. " new_options = list(lorebook_dropdown.options)\n",
  789. " removed_idx = new_options.index(removed_key)\n",
  790. " new_options.remove(removed_key)\n",
  791. " lorebook_dropdown.options = tuple(new_options)\n",
  792. " lorebook_dropdown.value = lorebook_dropdown.options[max(removed_idx-1, 0)]\n",
  793. "\n",
  794. "def lorebook_add_clicked(b):\n",
  795. " add_new_entry()\n",
  796. " update_buttons_visible()\n",
  797. " cur_state.clear_kv()\n",
  798. "\n",
  799. "def lorebook_remove_clicked(b):\n",
  800. " remove_entry()\n",
  801. " update_buttons_visible()\n",
  802. " cur_state.clear_kv()\n",
  803. "\n",
  804. "def lorebook_apply_changes_clicked(b):\n",
  805. " apply_changes_to_entry(True)\n",
  806. " update_buttons_visible()\n",
  807. " cur_state.clear_kv()\n",
  808. "\n",
  809. "def on_lorebook_dropdown_select(b):\n",
  810. " apply_changes_to_entry(False)\n",
  811. " if len(cur_state.lorebook_dict) > 0:\n",
  812. " lorebook_key_field.value = lorebook_dropdown.value\n",
  813. " lorebook_value_field.value = cur_state.lorebook_dict[lorebook_key_field.value]\n",
  814. " cur_state.lorebook_prev_key = lorebook_key_field.value\n",
  815. " cur_state.lorebook_prev_value = lorebook_value_field.value\n",
  816. " else:\n",
  817. " lorebook_key_field.value = ''\n",
  818. " lorebook_value_field.value = ''\n",
  819. " update_buttons_visible()\n",
  820. " cur_state.clear_kv()\n",
  821. "\n",
  822. "def on_update_lorebook_key_field(b):\n",
  823. " if lorebook_apply_changes_button.disabled:\n",
  824. " lorebook_apply_changes_button.disabled = False\n",
  825. "\n",
  826. "def on_update_lorebook_value_field(b):\n",
  827. " if lorebook_apply_changes_button.disabled:\n",
  828. " lorebook_apply_changes_button.disabled = False\n",
  829. "\n",
  830. "send_button.on_click(send_button_clicked)\n",
  831. "undo_button.on_click(undo_button_clicked)\n",
  832. "redo_button.on_click(redo_button_clicked)\n",
  833. "retry_button.on_click(retry_button_clicked)\n",
  834. "prev_retry_button.on_click(prev_retry_button_clicked)\n",
  835. "memory_button.observe(memory_button_clicked, names='value')\n",
  836. "context_button.observe(context_button_clicked, names='value')\n",
  837. "input_text_area.observe(on_update_input_text_area, names='value')\n",
  838. "lorebook_add_button.on_click(lorebook_add_clicked)\n",
  839. "lorebook_remove_button.on_click(lorebook_remove_clicked)\n",
  840. "lorebook_apply_changes_button.on_click(lorebook_apply_changes_clicked)\n",
  841. "lorebook_dropdown.observe(on_lorebook_dropdown_select, names='value')\n",
  842. "lorebook_key_field.observe(on_update_lorebook_key_field, names='value')\n",
  843. "lorebook_value_field.observe(on_update_lorebook_value_field, names='value')\n",
  844. "\n",
  845. "update_buttons_visible()\n",
  846. "\n",
  847. "display(main_panel, output)"
  848. ],
  849. "metadata": {
  850. "id": "JNrRMsUkoX99"
  851. },
  852. "execution_count": null,
  853. "outputs": []
  854. }
  855. ]
  856. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement
Public Pastes
Advertisement
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!