server.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588
  1. import re
  2. import gc
  3. import time
  4. import glob
  5. import torch
  6. import argparse
  7. import json
  8. from sys import exit
  9. from pathlib import Path
  10. import gradio as gr
  11. import warnings
  12. from tqdm import tqdm
  13. import transformers
  14. from transformers import AutoTokenizer, AutoModelForCausalLM
  15. from modules.html_generator import *
  16. from modules.ui import *
  17. from modules.stopping_criteria import _SentinelTokenStoppingCriteria
  18. transformers.logging.set_verbosity_error()
  19. parser = argparse.ArgumentParser()
  20. parser.add_argument('--model', type=str, help='Name of the model to load by default.')
  21. parser.add_argument('--notebook', action='store_true', help='Launch the web UI in notebook mode, where the output is written to the same text box as the input.')
  22. parser.add_argument('--chat', action='store_true', help='Launch the web UI in chat mode.')
  23. parser.add_argument('--cai-chat', action='store_true', help='Launch the web UI in chat mode with a style similar to Character.AI\'s. If the file profile.png or profile.jpg exists in the same folder as server.py, this image will be used as the bot\'s profile picture.')
  24. parser.add_argument('--cpu', action='store_true', help='Use the CPU to generate text.')
  25. parser.add_argument('--load-in-8bit', action='store_true', help='Load the model with 8-bit precision.')
  26. parser.add_argument('--auto-devices', action='store_true', help='Automatically split the model across the available GPU(s) and CPU.')
  27. parser.add_argument('--disk', action='store_true', help='If the model is too large for your GPU(s) and CPU combined, send the remaining layers to the disk.')
  28. parser.add_argument('--disk-cache-dir', type=str, help='Directory to save the disk cache to. Defaults to "cache/".')
  29. parser.add_argument('--gpu-memory', type=int, help='Maximum GPU memory in GiB to allocate. This is useful if you get out of memory errors while trying to generate text. Must be an integer number.')
  30. parser.add_argument('--cpu-memory', type=int, help='Maximum CPU memory in GiB to allocate for offloaded weights. Must be an integer number. Defaults to 99.')
  31. parser.add_argument('--no-stream', action='store_true', help='Don\'t stream the text output in real time. This improves the text generation performance.')
  32. parser.add_argument('--settings', type=str, help='Load the default interface settings from this json file. See settings-template.json for an example.')
  33. parser.add_argument('--listen', action='store_true', help='Make the web UI reachable from your local network.')
  34. parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.')
  35. args = parser.parse_args()
  36. if (args.chat or args.cai_chat) and not args.no_stream:
  37. print("Warning: chat mode currently becomes a lot slower with text streaming on.\nConsider starting the web UI with the --no-stream option.\n")
  38. settings = {
  39. 'max_new_tokens': 200,
  40. 'max_new_tokens_min': 1,
  41. 'max_new_tokens_max': 2000,
  42. 'preset': 'NovelAI-Sphinx Moth',
  43. 'name1': 'Person 1',
  44. 'name2': 'Person 2',
  45. 'context': 'This is a conversation between two people.',
  46. 'prompt': 'Common sense questions and answers\n\nQuestion: \nFactual answer:',
  47. 'prompt_gpt4chan': '-----\n--- 865467536\nInput text\n--- 865467537\n',
  48. 'stop_at_newline': True,
  49. 'history_size': 0,
  50. 'history_size_min': 0,
  51. 'history_size_max': 64,
  52. 'preset_pygmalion': 'Pygmalion',
  53. 'name1_pygmalion': 'You',
  54. 'name2_pygmalion': 'Kawaii',
  55. 'context_pygmalion': "Kawaii's persona: Kawaii is a cheerful person who loves to make others smile. She is an optimist who loves to spread happiness and positivity wherever she goes.\n<START>",
  56. 'stop_at_newline_pygmalion': False,
  57. }
  58. if args.settings is not None and Path(args.settings).exists():
  59. with open(Path(args.settings), 'r') as f:
  60. new_settings = json.load(f)
  61. for item in new_settings:
  62. if item in settings:
  63. settings[item] = new_settings[item]
  64. def load_model(model_name):
  65. print(f"Loading {model_name}...")
  66. t0 = time.time()
  67. # Default settings
  68. if not (args.cpu or args.load_in_8bit or args.auto_devices or args.disk or args.gpu_memory is not None):
  69. if Path(f"torch-dumps/{model_name}.pt").exists():
  70. print("Loading in .pt format...")
  71. model = torch.load(Path(f"torch-dumps/{model_name}.pt"))
  72. elif model_name.lower().startswith(('gpt-neo', 'opt-', 'galactica')) and any(size in model_name.lower() for size in ('13b', '20b', '30b')):
  73. model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), device_map='auto', load_in_8bit=True)
  74. else:
  75. model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
  76. # Custom
  77. else:
  78. settings = ["low_cpu_mem_usage=True"]
  79. command = "AutoModelForCausalLM.from_pretrained"
  80. if args.cpu:
  81. settings.append("torch_dtype=torch.float32")
  82. else:
  83. settings.append("device_map='auto'")
  84. if args.gpu_memory is not None:
  85. if args.cpu_memory is not None:
  86. settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '{args.cpu_memory}GiB'}}")
  87. else:
  88. settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '99GiB'}}")
  89. if args.disk:
  90. if args.disk_cache_dir is not None:
  91. settings.append(f"offload_folder='{args.disk_cache_dir}'")
  92. else:
  93. settings.append("offload_folder='cache'")
  94. if args.load_in_8bit:
  95. settings.append("load_in_8bit=True")
  96. else:
  97. settings.append("torch_dtype=torch.float16")
  98. settings = ', '.join(set(settings))
  99. command = f"{command}(Path(f'models/{model_name}'), {settings})"
  100. model = eval(command)
  101. # Loading the tokenizer
  102. if model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')) and Path(f"models/gpt-j-6B/").exists():
  103. tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/"))
  104. else:
  105. tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{model_name}/"))
  106. tokenizer.truncation_side = 'left'
  107. print(f"Loaded the model in {(time.time()-t0):.2f} seconds.")
  108. return model, tokenizer
  109. # Removes empty replies from gpt4chan outputs
  110. def fix_gpt4chan(s):
  111. for i in range(10):
  112. s = re.sub("--- [0-9]*\n>>[0-9]*\n---", "---", s)
  113. s = re.sub("--- [0-9]*\n *\n---", "---", s)
  114. s = re.sub("--- [0-9]*\n\n\n---", "---", s)
  115. return s
  116. # Fix the LaTeX equations in galactica
  117. def fix_galactica(s):
  118. s = s.replace(r'\[', r'$')
  119. s = s.replace(r'\]', r'$')
  120. s = s.replace(r'\(', r'$')
  121. s = s.replace(r'\)', r'$')
  122. s = s.replace(r'$$', r'$')
  123. return s
  124. def encode(prompt, tokens_to_generate=0, add_special_tokens=True):
  125. if args.cpu:
  126. input_ids = tokenizer.encode(str(prompt), return_tensors='pt', truncation=True, max_length=2048-tokens_to_generate, add_special_tokens=add_special_tokens)
  127. else:
  128. torch.cuda.empty_cache()
  129. input_ids = tokenizer.encode(str(prompt), return_tensors='pt', truncation=True, max_length=2048-tokens_to_generate, add_special_tokens=add_special_tokens).cuda()
  130. return input_ids
  131. def decode(output_ids):
  132. reply = tokenizer.decode(output_ids, skip_special_tokens=True)
  133. reply = reply.replace(r'<|endoftext|>', '')
  134. return reply
  135. def formatted_outputs(reply, model_name):
  136. if not (args.chat or args.cai_chat):
  137. if model_name.lower().startswith('galactica'):
  138. reply = fix_galactica(reply)
  139. return reply, reply, generate_basic_html(reply)
  140. elif model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')):
  141. reply = fix_gpt4chan(reply)
  142. return reply, 'Only applicable for GALACTICA models.', generate_4chan_html(reply)
  143. else:
  144. return reply, 'Only applicable for GALACTICA models.', generate_basic_html(reply)
  145. else:
  146. return reply
  147. def generate_reply(question, tokens, inference_settings, selected_model, eos_token=None, stopping_string=None):
  148. global model, tokenizer, model_name, loaded_preset, preset
  149. if selected_model != model_name:
  150. model_name = selected_model
  151. model = tokenizer = None
  152. if not args.cpu:
  153. gc.collect()
  154. torch.cuda.empty_cache()
  155. model, tokenizer = load_model(model_name)
  156. if inference_settings != loaded_preset:
  157. with open(Path(f'presets/{inference_settings}.txt'), 'r') as infile:
  158. preset = infile.read()
  159. loaded_preset = inference_settings
  160. cuda = "" if args.cpu else ".cuda()"
  161. n = None if eos_token is None else tokenizer.encode(eos_token, return_tensors='pt')[0][-1]
  162. input_ids = encode(question, tokens)
  163. # The stopping_criteria code below was copied from
  164. # https://github.com/PygmalionAI/gradio-ui/blob/master/src/model.py
  165. if stopping_string is not None:
  166. t = encode(stopping_string, 0, add_special_tokens=False)
  167. stopping_criteria_list = transformers.StoppingCriteriaList([
  168. _SentinelTokenStoppingCriteria(
  169. sentinel_token_ids=t,
  170. starting_idx=len(input_ids[0]))
  171. ])
  172. else:
  173. stopping_criteria_list = None
  174. # Generate the entire reply at once
  175. if args.no_stream:
  176. t0 = time.time()
  177. output = eval(f"model.generate(input_ids, eos_token_id={n}, stopping_criteria=stopping_criteria_list, {preset}){cuda}")
  178. reply = decode(output[0])
  179. t1 = time.time()
  180. print(f"Output generated in {(t1-t0):.2f} seconds ({(len(output[0])-len(input_ids[0]))/(t1-t0):.2f} it/s)")
  181. yield formatted_outputs(reply, model_name)
  182. # Generate the reply 1 token at a time
  183. else:
  184. yield formatted_outputs(question, model_name)
  185. preset = preset.replace('max_new_tokens=tokens', 'max_new_tokens=8')
  186. for i in tqdm(range(tokens//8+1)):
  187. output = eval(f"model.generate(input_ids, eos_token_id={n}, stopping_criteria=stopping_criteria_list, {preset}){cuda}")
  188. reply = decode(output[0])
  189. if eos_token is not None and reply[-1] == eos_token:
  190. break
  191. yield formatted_outputs(reply, model_name)
  192. input_ids = output
  193. def get_available_models():
  194. return sorted(set([item.replace('.pt', '') for item in map(lambda x : str(x.name), list(Path('models/').glob('*'))+list(Path('torch-dumps/').glob('*'))) if not item.endswith('.txt')]), key=str.lower)
  195. def get_available_presets():
  196. return sorted(set(map(lambda x : '.'.join(str(x.name).split('.')[:-1]), Path('presets').glob('*.txt'))), key=str.lower)
  197. def get_available_characters():
  198. return ["None"] + sorted(set(map(lambda x : '.'.join(str(x.name).split('.')[:-1]), Path('characters').glob('*.json'))), key=str.lower)
  199. available_models = get_available_models()
  200. available_presets = get_available_presets()
  201. available_characters = get_available_characters()
  202. # Choosing the default model
  203. if args.model is not None:
  204. model_name = args.model
  205. else:
  206. if len(available_models) == 0:
  207. print("No models are available! Please download at least one.")
  208. exit(0)
  209. elif len(available_models) == 1:
  210. i = 0
  211. else:
  212. print("The following models are available:\n")
  213. for i,model in enumerate(available_models):
  214. print(f"{i+1}. {model}")
  215. print(f"\nWhich one do you want to load? 1-{len(available_models)}\n")
  216. i = int(input())-1
  217. print()
  218. model_name = available_models[i]
  219. model, tokenizer = load_model(model_name)
  220. loaded_preset = None
  221. # UI settings
  222. default_text = settings['prompt_gpt4chan'] if model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')) else settings['prompt']
  223. description = f"\n\n# Text generation lab\nGenerate text using Large Language Models.\n"
  224. css = ".my-4 {margin-top: 0} .py-6 {padding-top: 2.5rem} #refresh-button {flex: none; margin: 0; padding: 0; min-width: 50px; border: none; box-shadow: none; border-radius: 0} #download-label, #upload-label {min-height: 0}"
  225. if args.chat or args.cai_chat:
  226. history = []
  227. character = None
  228. # This gets the new line characters right.
  229. def clean_chat_message(text):
  230. text = text.replace('\n', '\n\n')
  231. text = re.sub(r"\n{3,}", "\n\n", text)
  232. text = text.strip()
  233. return text
  234. def generate_chat_prompt(text, tokens, name1, name2, context, history_size):
  235. text = clean_chat_message(text)
  236. rows = [f"{context.strip()}\n"]
  237. i = len(history)-1
  238. count = 0
  239. while i >= 0 and len(encode(''.join(rows), tokens)[0]) < 2048-tokens:
  240. rows.insert(1, f"{name2}: {history[i][1].strip()}\n")
  241. count += 1
  242. if not (history[i][0] == '<|BEGIN-VISIBLE-CHAT|>'):
  243. rows.insert(1, f"{name1}: {history[i][0].strip()}\n")
  244. count += 1
  245. i -= 1
  246. if history_size != 0 and count >= history_size:
  247. break
  248. rows.append(f"{name1}: {text}\n")
  249. rows.append(f"{name2}:")
  250. while len(rows) > 3 and len(encode(''.join(rows), tokens)[0]) >= 2048-tokens:
  251. rows.pop(1)
  252. rows.pop(1)
  253. question = ''.join(rows)
  254. return question
  255. def remove_example_dialogue_from_history(history):
  256. _history = copy.deepcopy(history)
  257. for i in range(len(_history)):
  258. if '<|BEGIN-VISIBLE-CHAT|>' in _history[i][0]:
  259. _history[i][0] = _history[i][0].replace('<|BEGIN-VISIBLE-CHAT|>', '')
  260. _history = _history[i:]
  261. break
  262. return _history
  263. def chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size):
  264. question = generate_chat_prompt(text, tokens, name1, name2, context, history_size)
  265. history.append(['', ''])
  266. eos_token = '\n' if check else None
  267. for reply in generate_reply(question, tokens, inference_settings, selected_model, eos_token=eos_token, stopping_string=f"\n{name1}:"):
  268. next_character_found = False
  269. previous_idx = [m.start() for m in re.finditer(f"(^|\n){name2}:", question)]
  270. idx = [m.start() for m in re.finditer(f"(^|\n){name2}:", reply)]
  271. idx = idx[len(previous_idx)-1]
  272. reply = reply[idx + len(f"\n{name2}:"):]
  273. if check:
  274. reply = reply.split('\n')[0].strip()
  275. else:
  276. idx = reply.find(f"\n{name1}:")
  277. if idx != -1:
  278. reply = reply[:idx]
  279. next_character_found = True
  280. reply = clean_chat_message(reply)
  281. history[-1] = [text, reply]
  282. if next_character_found:
  283. break
  284. # Prevent the chat log from flashing if something like "\nYo" is generated just
  285. # before "\nYou:" is completed
  286. tmp = f"\n{name1}:"
  287. next_character_substring_found = False
  288. for j in range(1, len(tmp)):
  289. if reply[-j:] == tmp[:j]:
  290. next_character_substring_found = True
  291. if not next_character_substring_found:
  292. yield remove_example_dialogue_from_history(history)
  293. yield remove_example_dialogue_from_history(history)
  294. def cai_chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size):
  295. for history in chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size):
  296. yield generate_chat_html(history, name1, name2, character)
  297. def regenerate_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size):
  298. last = history.pop()
  299. text = last[0]
  300. if args.cai_chat:
  301. for i in cai_chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size):
  302. yield i
  303. else:
  304. for i in chatbot_wrapper(text, tokens, inference_settings, selected_model, name1, name2, context, check, history_size):
  305. yield i
  306. def remove_last_message(name1, name2):
  307. last = history.pop()
  308. _history = remove_example_dialogue_from_history(history)
  309. if args.cai_chat:
  310. return generate_chat_html(_history, name1, name2, character), last[0]
  311. else:
  312. return _history, last[0]
  313. def clear_html():
  314. return generate_chat_html([], "", "", character)
  315. def clear_chat_log(_character, name1, name2):
  316. global history
  317. if _character != 'None':
  318. load_character(_character, name1, name2)
  319. else:
  320. history = []
  321. _history = remove_example_dialogue_from_history(history)
  322. if args.cai_chat:
  323. return generate_chat_html(_history, name1, name2, character)
  324. else:
  325. return _history
  326. def redraw_html(name1, name2):
  327. global history
  328. _history = remove_example_dialogue_from_history(history)
  329. return generate_chat_html(_history, name1, name2, character)
  330. def tokenize_dialogue(dialogue, name1, name2):
  331. dialogue = re.sub('<START>', '', dialogue)
  332. dialogue = re.sub('(\n|^)[Aa]non:', '\\1You:', dialogue)
  333. idx = [m.start() for m in re.finditer(f"(^|\n)({name1}|{name2}):", dialogue)]
  334. messages = []
  335. for i in range(len(idx)-1):
  336. messages.append(dialogue[idx[i]:idx[i+1]].strip())
  337. messages.append(dialogue[idx[-1]:].strip())
  338. history = []
  339. entry = ['', '']
  340. for i in messages:
  341. if i.startswith(f'{name1}:'):
  342. entry[0] = i[len(f'{name1}:'):].strip()
  343. elif i.startswith(f'{name2}:'):
  344. entry[1] = i[len(f'{name2}:'):].strip()
  345. if not (len(entry[0]) == 0 and len(entry[1]) == 0):
  346. history.append(entry)
  347. entry = ['', '']
  348. return history
  349. def save_history():
  350. if not Path('logs').exists():
  351. Path('logs').mkdir()
  352. with open(Path('logs/conversation.json'), 'w') as f:
  353. f.write(json.dumps({'data': history}, indent=2))
  354. return Path('logs/conversation.json')
  355. def load_history(file, name1, name2):
  356. global history
  357. file = file.decode('utf-8')
  358. try:
  359. j = json.loads(file)
  360. if 'data' in j:
  361. history = j['data']
  362. # Compatibility with Pygmalion AI's official web UI
  363. elif 'chat' in j:
  364. history = [':'.join(x.split(':')[1:]).strip() for x in j['chat']]
  365. if len(j['chat']) > 0 and j['chat'][0].startswith(f'{name2}:'):
  366. history = [['<|BEGIN-VISIBLE-CHAT|>', history[0]]] + [[history[i], history[i+1]] for i in range(1, len(history)-1, 2)]
  367. else:
  368. history = [[history[i], history[i+1]] for i in range(0, len(history)-1, 2)]
  369. except:
  370. history = tokenize_dialogue(file, name1, name2)
  371. def load_character(_character, name1, name2):
  372. global history, character
  373. context = ""
  374. history = []
  375. if _character != 'None':
  376. character = _character
  377. with open(Path(f'characters/{_character}.json'), 'r') as f:
  378. data = json.loads(f.read())
  379. name2 = data['char_name']
  380. if 'char_persona' in data and data['char_persona'] != '':
  381. context += f"{data['char_name']}'s Persona: {data['char_persona']}\n"
  382. if 'world_scenario' in data and data['world_scenario'] != '':
  383. context += f"Scenario: {data['world_scenario']}\n"
  384. context = f"{context.strip()}\n<START>\n"
  385. if 'example_dialogue' in data and data['example_dialogue'] != '':
  386. history = tokenize_dialogue(data['example_dialogue'], name1, name2)
  387. if 'char_greeting' in data and len(data['char_greeting'].strip()) > 0:
  388. history += [['<|BEGIN-VISIBLE-CHAT|>', data['char_greeting']]]
  389. else:
  390. history += [['<|BEGIN-VISIBLE-CHAT|>', "Hello there!"]]
  391. else:
  392. character = None
  393. context = settings['context_pygmalion']
  394. name2 = settings['name2_pygmalion']
  395. _history = remove_example_dialogue_from_history(history)
  396. if args.cai_chat:
  397. return name2, context, generate_chat_html(_history, name1, name2, character)
  398. else:
  399. return name2, context, _history
  400. suffix = '_pygmalion' if 'pygmalion' in model_name.lower() else ''
  401. with gr.Blocks(css=css+".h-\[40vh\] {height: 66.67vh} .gradio-container {max-width: 800px; margin-left: auto; margin-right: auto}", analytics_enabled=False) as interface:
  402. if args.cai_chat:
  403. display1 = gr.HTML(value=generate_chat_html([], "", "", character))
  404. else:
  405. display1 = gr.Chatbot()
  406. textbox = gr.Textbox(label='Input')
  407. btn = gr.Button("Generate")
  408. with gr.Row():
  409. stop = gr.Button("Stop")
  410. btn_regenerate = gr.Button("Regenerate")
  411. btn_remove_last = gr.Button("Remove last")
  412. btn_clear = gr.Button("Clear history")
  413. with gr.Row():
  414. with gr.Column():
  415. length_slider = gr.Slider(minimum=settings['max_new_tokens_min'], maximum=settings['max_new_tokens_max'], step=1, label='max_new_tokens', value=settings['max_new_tokens'])
  416. with gr.Row():
  417. model_menu = gr.Dropdown(choices=available_models, value=model_name, label='Model')
  418. create_refresh_button(model_menu, lambda : None, lambda : {"choices": get_available_models()}, "refresh-button")
  419. with gr.Column():
  420. history_size_slider = gr.Slider(minimum=settings['history_size_min'], maximum=settings['history_size_max'], step=1, label='Chat history size in prompt (0 for no limit)', value=settings['history_size'])
  421. with gr.Row():
  422. preset_menu = gr.Dropdown(choices=available_presets, value=settings[f'preset{suffix}'], label='Generation parameters preset')
  423. create_refresh_button(preset_menu, lambda : None, lambda : {"choices": get_available_presets()}, "refresh-button")
  424. name1 = gr.Textbox(value=settings[f'name1{suffix}'], lines=1, label='Your name')
  425. name2 = gr.Textbox(value=settings[f'name2{suffix}'], lines=1, label='Bot\'s name')
  426. context = gr.Textbox(value=settings[f'context{suffix}'], lines=2, label='Context')
  427. with gr.Row():
  428. character_menu = gr.Dropdown(choices=available_characters, value="None", label='Character')
  429. create_refresh_button(character_menu, lambda : None, lambda : {"choices": get_available_characters()}, "refresh-button")
  430. with gr.Row():
  431. check = gr.Checkbox(value=settings[f'stop_at_newline{suffix}'], label='Stop generating at new line character?')
  432. with gr.Row():
  433. with gr.Tab('Download chat history'):
  434. download = gr.File()
  435. save_btn = gr.Button(value="Click me")
  436. with gr.Tab('Upload chat history'):
  437. upload = gr.File(type='binary')
  438. input_params = [textbox, length_slider, preset_menu, model_menu, name1, name2, context, check, history_size_slider]
  439. if args.cai_chat:
  440. gen_event = btn.click(cai_chatbot_wrapper, input_params, display1, show_progress=args.no_stream, api_name="textgen")
  441. gen_event2 = textbox.submit(cai_chatbot_wrapper, input_params, display1, show_progress=args.no_stream)
  442. else:
  443. gen_event = btn.click(chatbot_wrapper, input_params, display1, show_progress=args.no_stream, api_name="textgen")
  444. gen_event2 = textbox.submit(chatbot_wrapper, input_params, display1, show_progress=args.no_stream)
  445. gen_event3 = btn_regenerate.click(regenerate_wrapper, input_params, display1, show_progress=args.no_stream)
  446. btn_clear.click(clear_chat_log, [character_menu, name1, name2], display1)
  447. btn_remove_last.click(remove_last_message, [name1, name2], [display1, textbox], show_progress=False)
  448. btn.click(lambda x: "", textbox, textbox, show_progress=False)
  449. btn_regenerate.click(lambda x: "", textbox, textbox, show_progress=False)
  450. textbox.submit(lambda x: "", textbox, textbox, show_progress=False)
  451. stop.click(None, None, None, cancels=[gen_event, gen_event2, gen_event3])
  452. save_btn.click(save_history, inputs=[], outputs=[download])
  453. upload.upload(load_history, [upload, name1, name2], [])
  454. character_menu.change(load_character, [character_menu, name1, name2], [name2, context, display1])
  455. if args.cai_chat:
  456. upload.upload(redraw_html, [name1, name2], [display1])
  457. else:
  458. upload.upload(lambda : remove_example_dialogue_from_history(history), [], [display1])
  459. elif args.notebook:
  460. with gr.Blocks(css=css, analytics_enabled=False) as interface:
  461. gr.Markdown(description)
  462. with gr.Tab('Raw'):
  463. textbox = gr.Textbox(value=default_text, lines=23)
  464. with gr.Tab('Markdown'):
  465. markdown = gr.Markdown()
  466. with gr.Tab('HTML'):
  467. html = gr.HTML()
  468. btn = gr.Button("Generate")
  469. stop = gr.Button("Stop")
  470. length_slider = gr.Slider(minimum=settings['max_new_tokens_min'], maximum=settings['max_new_tokens_max'], step=1, label='max_new_tokens', value=settings['max_new_tokens'])
  471. with gr.Row():
  472. with gr.Column():
  473. with gr.Row():
  474. model_menu = gr.Dropdown(choices=available_models, value=model_name, label='Model')
  475. create_refresh_button(model_menu, lambda : None, lambda : {"choices": get_available_models()}, "refresh-button")
  476. with gr.Column():
  477. with gr.Row():
  478. preset_menu = gr.Dropdown(choices=available_presets, value=settings['preset'], label='Generation parameters preset')
  479. create_refresh_button(preset_menu, lambda : None, lambda : {"choices": get_available_presets()}, "refresh-button")
  480. gen_event = btn.click(generate_reply, [textbox, length_slider, preset_menu, model_menu], [textbox, markdown, html], show_progress=args.no_stream, api_name="textgen")
  481. gen_event2 = textbox.submit(generate_reply, [textbox, length_slider, preset_menu, model_menu], [textbox, markdown, html], show_progress=args.no_stream)
  482. stop.click(None, None, None, cancels=[gen_event, gen_event2])
  483. else:
  484. with gr.Blocks(css=css, analytics_enabled=False) as interface:
  485. gr.Markdown(description)
  486. with gr.Row():
  487. with gr.Column():
  488. textbox = gr.Textbox(value=default_text, lines=15, label='Input')
  489. length_slider = gr.Slider(minimum=settings['max_new_tokens_min'], maximum=settings['max_new_tokens_max'], step=1, label='max_new_tokens', value=settings['max_new_tokens'])
  490. with gr.Row():
  491. preset_menu = gr.Dropdown(choices=available_presets, value=settings['preset'], label='Generation parameters preset')
  492. create_refresh_button(preset_menu, lambda : None, lambda : {"choices": get_available_presets()}, "refresh-button")
  493. with gr.Row():
  494. model_menu = gr.Dropdown(choices=available_models, value=model_name, label='Model')
  495. create_refresh_button(model_menu, lambda : None, lambda : {"choices": get_available_models()}, "refresh-button")
  496. btn = gr.Button("Generate")
  497. with gr.Row():
  498. with gr.Column():
  499. cont = gr.Button("Continue")
  500. with gr.Column():
  501. stop = gr.Button("Stop")
  502. with gr.Column():
  503. with gr.Tab('Raw'):
  504. output_textbox = gr.Textbox(lines=15, label='Output')
  505. with gr.Tab('Markdown'):
  506. markdown = gr.Markdown()
  507. with gr.Tab('HTML'):
  508. html = gr.HTML()
  509. gen_event = btn.click(generate_reply, [textbox, length_slider, preset_menu, model_menu], [output_textbox, markdown, html], show_progress=args.no_stream, api_name="textgen")
  510. gen_event2 = textbox.submit(generate_reply, [textbox, length_slider, preset_menu, model_menu], [output_textbox, markdown, html], show_progress=args.no_stream)
  511. cont_event = cont.click(generate_reply, [output_textbox, length_slider, preset_menu, model_menu], [output_textbox, markdown, html], show_progress=args.no_stream)
  512. stop.click(None, None, None, cancels=[gen_event, gen_event2, cont_event])
  513. interface.queue()
  514. if args.listen:
  515. interface.launch(share=args.share, server_name="0.0.0.0")
  516. else:
  517. interface.launch(share=args.share)