m-ric/chunk_visualizer

1import gradio as gr2from langchain.text_splitter import (3    CharacterTextSplitter,4    RecursiveCharacterTextSplitter,5    Language,6)7from transformers import AutoTokenizer8from overlap import unoverlap_list9 10LABEL_TEXTSPLITTER = "🦜🔗 LangChain's CharacterTextSplitter"11LABEL_RECURSIVE = "🦜🔗 LangChain's RecursiveCharacterTextSplitter"12 13bert_tokenizer = AutoTokenizer.from_pretrained('google-bert/bert-base-uncased')14 15def length_tokens(txt):16    return len(bert_tokenizer.tokenize(txt))17 18 19def extract_separators_from_string(separators_str):20    try:21        separators_str = separators_str.replace("\\n", "\n").replace("\\t", "\t").replace("\\\\", "\\") # fix special characters22        separators = separators_str[1:-1].split(", ")23        return [separator.replace('"', "").replace("'", "") for separator in separators]24    except Exception as e:25        raise gr.Error(f"""26        Did not succeed in extracting seperators from string: {separator_str} due to: {str(e)}.27        Please type it in the correct format: "['separator_1', 'separator_2', ...]"28        """)29 30def change_split_selection(split_selection):31    return (32        gr.Textbox.update(visible=(split_selection==LABEL_RECURSIVE)),33        gr.Radio.update(visible=(split_selection==LABEL_RECURSIVE)),34    )35 36def chunk(text, length, splitter_selection, separators_str, length_unit_selection, chunk_overlap):37    separators = extract_separators_from_string(separators_str)38    length_function = (length_tokens if "token" in length_unit_selection.lower() else len)39    if splitter_selection == LABEL_TEXTSPLITTER:40        text_splitter = CharacterTextSplitter(41            chunk_size=length,42            chunk_overlap=int(chunk_overlap),43            length_function=length_function,44            strip_whitespace=False,45            is_separator_regex=False,46            separator=" ",47        )48    elif splitter_selection == LABEL_RECURSIVE:49        text_splitter = RecursiveCharacterTextSplitter(50            chunk_size=length,51            chunk_overlap=int(chunk_overlap),52            length_function=length_function,53            strip_whitespace=False,54            separators=separators,55        )56    splits = text_splitter.create_documents([text])57    text_splits = [split.page_content for split in splits]58    unoverlapped_text_splits = unoverlap_list(text_splits)59    output = [((split[0], 'Overlap') if split[1] else (split[0], f"Chunk {str(i)}")) for i, split in enumerate(unoverlapped_text_splits)]60    return output61 62def change_preset_separators(choice):63    text_splitter = RecursiveCharacterTextSplitter()64    if choice == "Default":65        return ["\n\n", "\n", " ", ""]66    elif choice == "Markdown":67        return text_splitter.get_separators_for_language(Language.MARKDOWN)68    elif choice == "Python":69        return text_splitter.get_separators_for_language(Language.PYTHON)70    else:71        raise gr.Error("Choice of preset not recognized.")72 73 74EXAMPLE_TEXT = """### Chapter 675 76WHAT SORT OF DESPOTISM DEMOCRATIC NATIONS HAVE TO FEAR77 78I had remarked during my stay in the United States that a democratic state of society, similar to that of the Americans, might offer singular facilities for the establishment of despotism; and I perceived, upon my return to Europe, how much use had already been made, by most of our rulers, of the notions, the sentiments, and the wants created by this same social condition, for the purpose of extending the circle of their power. This led me to think that the nations of Christendom would perhaps eventually undergo some oppression like that which hung over several of the nations of the ancient world.79A more accurate examination of the subject, and five years of further meditation, have not diminished my fears, but have changed their object.80No sovereign ever lived in former ages so absolute or so powerful as to undertake to administer by his own agency, and without the assistance of intermediate powers, all the parts of a great empire; none ever attempted to subject all his subjects indiscriminately to strict uniformity of regulation and personally to tutor and direct every member of the community. The notion of such an undertaking never occurred to the human mind; and if any man had conceived it, the want of information, the imperfection of the administrative system, and, above all, the natural obstacles caused by the inequality of conditions would speedily have checked the execution of so vast a design.81 82---83 84### Challenges of agent systems85 86Generally, the difficult parts of running an agent system for the LLM engine are:87 881. From supplied tools, choose the one that will help advance to a desired goal: e.g. when asked `"What is the smallest prime number greater than 30,000?"`, the agent could call the `Search` tool with `"What is he height of K2"` but it won't help.892. Call tools with a rigorous argument formatting: for instance when trying to calculate the speed of a car that went 3 km in 10 minutes, you have to call tool `Calculator` to divide `distance` by `time` : even if your Calculator tool accepts calls in the JSON format: `{”tool”: “Calculator”, “args”: “3km/10min”}` , there are many pitfalls, for instance:90    - Misspelling the tool name: `“calculator”` or `“Compute”` wouldn’t work91    - Giving the name of the arguments instead of their values: `“args”: “distance/time”`92    - Non-standardized formatting: `“args": "3km in 10minutes”`933. Efficiently ingesting and using the information gathered in the past observations, be it the initial context or the observations returned after using tool uses.94 95 96So, how would a complete Agent setup look like?97 98## Running agents with LangChain99 100We have just integrated a `ChatHuggingFace` wrapper that lets you create agents based on open-source models in [🦜🔗LangChain](https://www.langchain.com/).101 102The code to create the ChatModel and give it tools is really simple, you can check it all in the [Langchain doc](https://python.langchain.com/docs/integrations/chat/huggingface). 103 104```python105from langchain_community.llms import HuggingFaceHub106from langchain_community.chat_models.huggingface import ChatHuggingFace107 108llm = HuggingFaceHub(109    repo_id="HuggingFaceH4/zephyr-7b-beta",110    task="text-generation",111)112 113chat_model = ChatHuggingFace(llm=llm)114```115"""116 117    118with gr.Blocks(theme=gr.themes.Citrus(text_size='md', font=["monospace"], primary_hue=gr.themes.colors.green)) as demo:119    text = gr.Textbox(label="Your text 🪶", value=EXAMPLE_TEXT)120    with gr.Row():121        split_selection = gr.Dropdown(122            choices=[123                LABEL_TEXTSPLITTER,124                LABEL_RECURSIVE,125            ],126            value=LABEL_RECURSIVE,127            label="Method to split chunks 🍞",128        )129        separators_selection = gr.Textbox(130            elem_id="textbox_id",131            value=["\n\n", "\n", " ", ""],132            info="Separators used in RecursiveCharacterTextSplitter",133            show_label=False, # or set label to an empty string if you want to keep its space134            visible=True,135        )136        separator_preset_selection = gr.Radio(137            ['Default', 'Python', 'Markdown'],138            label="Choose a preset",139            info="This will apply a specific set of separators to RecursiveCharacterTextSplitter.",140            visible=True,141        )142    with gr.Row():143        length_unit_selection = gr.Dropdown(144            choices=[145                "Character count",146                "Token count (BERT tokens)",147            ],148            value="Character count",149            label="Length function",150            info="How should we measure our chunk lengths?",151        )152        slider_count = gr.Slider(153            50, 500, value=200, step=1, label="Chunk length 📏", info="In the chosen unit."154        )155        chunk_overlap = gr.Slider(156            0, 50, value=10, step=1, label="Overlap between chunks", info="In the chosen unit."157        )158    out = gr.HighlightedText(159        label="Output",160        show_legend=True,161        show_label=False,162        color_map={'Overlap': '#DADADA'}163    )164 165    split_selection.change(166        fn=change_split_selection,167        inputs=split_selection,168        outputs=[separators_selection, separator_preset_selection],169    )170    separator_preset_selection.change(171        fn=change_preset_separators,172        inputs=separator_preset_selection,173        outputs=separators_selection,174    )175    gr.on(176        [text.change, length_unit_selection.change, separators_selection.change, split_selection.change, slider_count.change, chunk_overlap.change],177        chunk,178        [text, slider_count, split_selection, separators_selection, length_unit_selection, chunk_overlap],179        outputs=out180    )181    demo.load(chunk, inputs=[text, slider_count, split_selection, separators_selection, length_unit_selection, chunk_overlap], outputs=out)182demo.launch()