| """ |
| An example of extension. It does nothing, but you can add transformations |
| before the return statements to customize the webui behavior. |
| |
| Starting from history_modifier and ending in output_modifier, the |
| functions are declared in the same order that they are called at |
| generation time. |
| """ |
|
|
| import gradio as gr |
| import torch |
| from transformers import LogitsProcessor |
|
|
| from modules import chat, shared |
| from modules.text_generation import ( |
| decode, |
| encode, |
| generate_reply, |
| ) |
|
|
| params = { |
| "display_name": "Example Extension", |
| "is_tab": False, |
| } |
|
|
| class MyLogits(LogitsProcessor): |
| """ |
| Manipulates the probabilities for the next token before it gets sampled. |
| Used in the logits_processor_modifier function below. |
| """ |
| def __init__(self): |
| pass |
|
|
| def __call__(self, input_ids, scores): |
| |
| |
| |
| return scores |
|
|
| def history_modifier(history): |
| """ |
| Modifies the chat history. |
| Only used in chat mode. |
| """ |
| return history |
|
|
| def state_modifier(state): |
| """ |
| Modifies the state variable, which is a dictionary containing the input |
| values in the UI like sliders and checkboxes. |
| """ |
| return state |
|
|
| def chat_input_modifier(text, visible_text, state): |
| """ |
| Modifies the user input string in chat mode (visible_text). |
| You can also modify the internal representation of the user |
| input (text) to change how it will appear in the prompt. |
| """ |
| return text, visible_text |
|
|
| def input_modifier(string, state, is_chat=False): |
| """ |
| In default/notebook modes, modifies the whole prompt. |
| |
| In chat mode, it is the same as chat_input_modifier but only applied |
| to "text", here called "string", and not to "visible_text". |
| """ |
| return string |
|
|
| def bot_prefix_modifier(string, state): |
| """ |
| Modifies the prefix for the next bot reply in chat mode. |
| By default, the prefix will be something like "Bot Name:". |
| """ |
| return string |
|
|
| def tokenizer_modifier(state, prompt, input_ids, input_embeds): |
| """ |
| Modifies the input ids and embeds. |
| Used by the multimodal extension to put image embeddings in the prompt. |
| Only used by loaders that use the transformers library for sampling. |
| """ |
| return prompt, input_ids, input_embeds |
|
|
| def logits_processor_modifier(processor_list, input_ids): |
| """ |
| Adds logits processors to the list, allowing you to access and modify |
| the next token probabilities. |
| Only used by loaders that use the transformers library for sampling. |
| """ |
| processor_list.append(MyLogits()) |
| return processor_list |
|
|
| def output_modifier(string, state, is_chat=False): |
| """ |
| Modifies the LLM output before it gets presented. |
| |
| In chat mode, the modified version goes into history['visible'], |
| and the original version goes into history['internal']. |
| """ |
| return string |
|
|
| def custom_generate_chat_prompt(user_input, state, **kwargs): |
| """ |
| Replaces the function that generates the prompt from the chat history. |
| Only used in chat mode. |
| """ |
| result = chat.generate_chat_prompt(user_input, state, **kwargs) |
| return result |
|
|
| def custom_css(): |
| """ |
| Returns a CSS string that gets appended to the CSS for the webui. |
| """ |
| return '' |
|
|
| def custom_js(): |
| """ |
| Returns a javascript string that gets appended to the javascript |
| for the webui. |
| """ |
| return '' |
|
|
| def setup(): |
| """ |
| Gets executed only once, when the extension is imported. |
| """ |
| pass |
|
|
| def ui(): |
| """ |
| Gets executed when the UI is drawn. Custom gradio elements and |
| their corresponding event handlers should be defined here. |
| |
| To learn about gradio components, check out the docs: |
| https://gradio.app/docs/ |
| """ |
| pass |
|
|
