Lsavints/swo_llm
0
1from typing import Iterator2 3import gradio as gr4import torch5 6from model import get_input_token_length, run7 8DEFAULT_SYSTEM_PROMPT = """\9Play a role of customer support assistant in marketplace of software products of SoftwareOne company. Your task is to provide short but correct answers to the client regarding software products and services. 10If you don't know the answer to a question, don't share false information."""11MAX_MAX_NEW_TOKENS = 204812DEFAULT_MAX_NEW_TOKENS = 25613MAX_INPUT_TOKEN_LENGTH = 400014 15DESCRIPTION = """16This is a prototype of SWO LLM17"""18 19if not torch.cuda.is_available():20 DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'21 22 23def clear_and_save_textbox(message: str) -> tuple[str, str]:24 return '', message25 26 27def display_input(message: str,28 history: list[tuple[str, str]]) -> list[tuple[str, str]]:29 history.append((message, ''))30 return history31 32 33def delete_prev_fn(34 history: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], str]:35 try:36 message, _ = history.pop()37 except IndexError:38 message = ''39 return history, message or ''40 41 42def generate(43 message: str,44 history_with_input: list[tuple[str, str]],45 system_prompt: str,46 max_new_tokens: int,47 temperature: float,48 top_p: float,49 top_k: int,50) -> Iterator[list[tuple[str, str]]]:51 if max_new_tokens > MAX_MAX_NEW_TOKENS:52 raise ValueError53 54 history = history_with_input[:-1]55 generator = run(message, history, system_prompt, max_new_tokens, temperature, top_p, top_k)56 try:57 first_response = next(generator)58 yield history + [(message, first_response)]59 except StopIteration:60 yield history + [(message, '')]61 for response in generator:62 yield history + [(message, response)]63 64 65def process_example(message: str) -> tuple[str, list[tuple[str, str]]]:66 generator = generate(message, [], DEFAULT_SYSTEM_PROMPT, 1024, 1, 0.95, 50)67 for x in generator:68 pass69 return '', x70 71 72def check_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> None:73 input_token_length = get_input_token_length(message, chat_history, system_prompt)74 if input_token_length > MAX_INPUT_TOKEN_LENGTH:75 raise gr.Error(f'The accumulated input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Clear your chat history and try again.')76 77 78with gr.Blocks(css='style.css') as demo:79 gr.Markdown(DESCRIPTION)80 81 with gr.Group():82 chatbot = gr.Chatbot(label='Chatbot')83 with gr.Row():84 textbox = gr.Textbox(85 container=False,86 show_label=False,87 placeholder='Type a message...',88 scale=10,89 )90 submit_button = gr.Button('Submit',91 variant='primary',92 scale=1,93 min_width=0)94 with gr.Row():95 retry_button = gr.Button('🔄 Retry', variant='secondary')96 undo_button = gr.Button('↩️ Undo', variant='secondary')97 clear_button = gr.Button('🗑️ Clear', variant='secondary')98 99 saved_input = gr.State()100 101 with gr.Accordion(label='Advanced options', open=False):102 system_prompt = gr.Textbox(label='System prompt',103 value=DEFAULT_SYSTEM_PROMPT,104 lines=6)105 max_new_tokens = gr.Slider(106 label='Max new tokens',107 minimum=1,108 maximum=MAX_MAX_NEW_TOKENS,109 step=1,110 value=DEFAULT_MAX_NEW_TOKENS,111 )112 temperature = gr.Slider(113 label='Temperature',114 minimum=0.1,115 maximum=4.0,116 step=0.1,117 value=0.2,118 )119 top_p = gr.Slider(120 label='Top-p (nucleus sampling)',121 minimum=0.05,122 maximum=1.0,123 step=0.05,124 value=0.95,125 )126 top_k = gr.Slider(127 label='Top-k',128 minimum=1,129 maximum=1000,130 step=1,131 value=7,132 )133 134 gr.Examples(135 examples=[136 'Hello there! How are you doing?',137 'Can you explain briefly function of Client Portal?',138 'How many employees are in SoftwareOne?',139 ],140 inputs=textbox,141 outputs=[textbox, chatbot],142 fn=process_example,143 cache_examples=True,144 )145 146 textbox.submit(147 fn=clear_and_save_textbox,148 inputs=textbox,149 outputs=[textbox, saved_input],150 api_name=False,151 queue=False,152 ).then(153 fn=display_input,154 inputs=[saved_input, chatbot],155 outputs=chatbot,156 api_name=False,157 queue=False,158 ).then(159 fn=check_input_token_length,160 inputs=[saved_input, chatbot, system_prompt],161 api_name=False,162 queue=False,163 ).success(164 fn=generate,165 inputs=[166 saved_input,167 chatbot,168 system_prompt,169 max_new_tokens,170 temperature,171 top_p,172 top_k,173 ],174 outputs=chatbot,175 api_name=False,176 )177 178 button_event_preprocess = submit_button.click(179 fn=clear_and_save_textbox,180 inputs=textbox,181 outputs=[textbox, saved_input],182 api_name=False,183 queue=False,184 ).then(185 fn=display_input,186 inputs=[saved_input, chatbot],187 outputs=chatbot,188 api_name=False,189 queue=False,190 ).then(191 fn=check_input_token_length,192 inputs=[saved_input, chatbot, system_prompt],193 api_name=False,194 queue=False,195 ).success(196 fn=generate,197 inputs=[198 saved_input,199 chatbot,200 system_prompt,201 max_new_tokens,202 temperature,203 top_p,204 top_k,205 ],206 outputs=chatbot,207 api_name=False,208 )209 210 retry_button.click(211 fn=delete_prev_fn,212 inputs=chatbot,213 outputs=[chatbot, saved_input],214 api_name=False,215 queue=False,216 ).then(217 fn=display_input,218 inputs=[saved_input, chatbot],219 outputs=chatbot,220 api_name=False,221 queue=False,222 ).then(223 fn=generate,224 inputs=[225 saved_input,226 chatbot,227 system_prompt,228 max_new_tokens,229 temperature,230 top_p,231 top_k,232 ],233 outputs=chatbot,234 api_name=False,235 )236 237 undo_button.click(238 fn=delete_prev_fn,239 inputs=chatbot,240 outputs=[chatbot, saved_input],241 api_name=False,242 queue=False,243 ).then(244 fn=lambda x: x,245 inputs=[saved_input],246 outputs=textbox,247 api_name=False,248 queue=False,249 )250 251 clear_button.click(252 fn=lambda: ([], ''),253 outputs=[chatbot, saved_input],254 queue=False,255 api_name=False,256 )257 258demo.queue(max_size=20).launch()259