import os
import tempfile
from typing import List, Union

import streamlit as st
import tiktoken
from langchain.text_splitter import (
    CharacterTextSplitter,
    RecursiveCharacterTextSplitter,
)
from langchain.text_splitter import (
    TextSplitter as LCSplitter,
)
from langchain.text_splitter import TokenTextSplitter as LCTokenTextSplitter
from llama_index import SimpleDirectoryReader
from llama_index.node_parser.interface import TextSplitter
from llama_index.schema import Document
from llama_index.text_splitter import CodeSplitter, SentenceSplitter, TokenTextSplitter
from streamlit.runtime.uploaded_file_manager import UploadedFile

DEFAULT_TEXT = "The quick brown fox jumps over the lazy dog."

text = st.sidebar.text_area("Enter text", value=DEFAULT_TEXT)
uploaded_files = st.sidebar.file_uploader("Upload file", accept_multiple_files=True)
type = st.sidebar.radio("Document Type", options=["Text", "Code"])
n_cols = st.sidebar.number_input("Columns", value=2, min_value=1, max_value=3)
assert isinstance(n_cols, int)


@st.cache_resource(ttl=3600)
def load_document(uploaded_files: List[UploadedFile]) -> List[Document]:
    # Read documents
    temp_dir = tempfile.TemporaryDirectory()
    for file in uploaded_files:
        temp_filepath = os.path.join(temp_dir.name, file.name)
        with open(temp_filepath, "wb") as f:
            f.write(file.getvalue())

    reader = SimpleDirectoryReader(input_dir=temp_dir.name)
    return reader.load_data()


if uploaded_files:
    if text != DEFAULT_TEXT:
        st.warning("Text will be ignored when uploading files")
    docs = load_document(uploaded_files)
    text = "\n".join([doc.text for doc in docs])


chunk_size = st.slider(
    "Chunk Size",
    value=512,
    min_value=1,
    max_value=4096,
)
chunk_overlap = st.slider(
    "Chunk Overlap",
    value=0,
    min_value=0,
    max_value=4096,
)

cols = st.columns(n_cols)
for ind, col in enumerate(cols):
    if type == "Text":
        text_splitter_cls = col.selectbox(
            "Text Splitter",
            options=[
                "TokenTextSplitter",
                "SentenceSplitter",
                "LC:RecursiveCharacterTextSplitter",
                "LC:CharacterTextSplitter",
                "LC:TokenTextSplitter",
            ],
            index=ind,
            key=f"splitter_cls_{ind}",
        )

        text_splitter: Union[TextSplitter, LCSplitter]
        if text_splitter_cls == "TokenTextSplitter":
            text_splitter = TokenTextSplitter(
                chunk_size=chunk_size, chunk_overlap=chunk_overlap
            )
        elif text_splitter_cls == "SentenceSplitter":
            text_splitter = SentenceSplitter(
                chunk_size=chunk_size, chunk_overlap=chunk_overlap
            )
        elif text_splitter_cls == "LC:RecursiveCharacterTextSplitter":
            text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
                chunk_size=chunk_size, chunk_overlap=chunk_overlap
            )
        elif text_splitter_cls == "LC:CharacterTextSplitter":
            text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
                chunk_size=chunk_size, chunk_overlap=chunk_overlap
            )
        elif text_splitter_cls == "LC:TokenTextSplitter":
            text_splitter = LCTokenTextSplitter(
                chunk_size=chunk_size, chunk_overlap=chunk_overlap
            )
        else:
            raise ValueError("Unknown text splitter")
    elif type == "Code":
        text_splitter_cls = col.selectbox("Text Splitter", options=["CodeSplitter"])
        if text_splitter_cls == "CodeSplitter":
            language = col.text_input("Language", value="python")
            max_chars = col.slider("Max Chars", value=1500)

            text_splitter = CodeSplitter(language=language, max_chars=max_chars)
        else:
            raise ValueError("Unknown text splitter")

    chunks = text_splitter.split_text(text)
    tokenizer = tiktoken.get_encoding("gpt2").encode

    for chunk_ind, chunk in enumerate(chunks):
        n_tokens = len(tokenizer(chunk))
        n_chars = len(chunk)
        col.text_area(
            f"Chunk {chunk_ind} - {n_tokens} tokens - {n_chars} chars",
            chunk,
            key=f"text_area_{ind}_{chunk_ind}",
            height=500,
        )

langchain.text_splitter.RecursiveCharacterTextSplitter.from_tiktoken_encoder langchain.text_splitter.CharacterTextSplitter.from_tiktoken_encoder langchain.text_splitter.TokenTextSplitter
[((718, 772), 'streamlit.sidebar.text_area', 'st.sidebar.text_area', (['"""Enter text"""'], {'value': 'DEFAULT_TEXT'}), "('Enter text', value=DEFAULT_TEXT)\n", (738, 772), True, 'import streamlit as st\n'), ((790, 857), 'streamlit.sidebar.file_uploader', 'st.sidebar.file_uploader', (['"""Upload file"""'], {'accept_multiple_files': '(True)'}), "('Upload file', accept_multiple_files=True)\n", (814, 857), True, 'import streamlit as st\n'), ((865, 924), 'streamlit.sidebar.radio', 'st.sidebar.radio', (['"""Document Type"""'], {'options': "['Text', 'Code']"}), "('Document Type', options=['Text', 'Code'])\n", (881, 924), True, 'import streamlit as st\n'), ((934, 1003), 'streamlit.sidebar.number_input', 'st.sidebar.number_input', (['"""Columns"""'], {'value': '(2)', 'min_value': '(1)', 'max_value': '(3)'}), "('Columns', value=2, min_value=1, max_value=3)\n", (957, 1003), True, 'import streamlit as st\n'), ((1038, 1065), 'streamlit.cache_resource', 'st.cache_resource', ([], {'ttl': '(3600)'}), '(ttl=3600)\n', (1055, 1065), True, 'import streamlit as st\n'), ((1692, 1755), 'streamlit.slider', 'st.slider', (['"""Chunk Size"""'], {'value': '(512)', 'min_value': '(1)', 'max_value': '(4096)'}), "('Chunk Size', value=512, min_value=1, max_value=4096)\n", (1701, 1755), True, 'import streamlit as st\n'), ((1791, 1855), 'streamlit.slider', 'st.slider', (['"""Chunk Overlap"""'], {'value': '(0)', 'min_value': '(0)', 'max_value': '(4096)'}), "('Chunk Overlap', value=0, min_value=0, max_value=4096)\n", (1800, 1855), True, 'import streamlit as st\n'), ((1883, 1901), 'streamlit.columns', 'st.columns', (['n_cols'], {}), '(n_cols)\n', (1893, 1901), True, 'import streamlit as st\n'), ((1175, 1204), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ([], {}), '()\n', (1202, 1204), False, 'import tempfile\n'), ((1396, 1442), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', ([], {'input_dir': 'temp_dir.name'}), '(input_dir=temp_dir.name)\n', (1417, 1442), False, 'from llama_index import SimpleDirectoryReader\n'), ((1261, 1299), 'os.path.join', 'os.path.join', (['temp_dir.name', 'file.name'], {}), '(temp_dir.name, file.name)\n', (1273, 1299), False, 'import os\n'), ((1531, 1586), 'streamlit.warning', 'st.warning', (['"""Text will be ignored when uploading files"""'], {}), "('Text will be ignored when uploading files')\n", (1541, 1586), True, 'import streamlit as st\n'), ((3968, 3997), 'tiktoken.get_encoding', 'tiktoken.get_encoding', (['"""gpt2"""'], {}), "('gpt2')\n", (3989, 3997), False, 'import tiktoken\n'), ((2486, 2555), 'llama_index.text_splitter.TokenTextSplitter', 'TokenTextSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n', (2503, 2555), False, 'from llama_index.text_splitter import CodeSplitter, SentenceSplitter, TokenTextSplitter\n'), ((2668, 2736), 'llama_index.text_splitter.SentenceSplitter', 'SentenceSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n', (2684, 2736), False, 'from llama_index.text_splitter import CodeSplitter, SentenceSplitter, TokenTextSplitter\n'), ((3786, 3838), 'llama_index.text_splitter.CodeSplitter', 'CodeSplitter', ([], {'language': 'language', 'max_chars': 'max_chars'}), '(language=language, max_chars=max_chars)\n', (3798, 3838), False, 'from llama_index.text_splitter import CodeSplitter, SentenceSplitter, TokenTextSplitter\n'), ((2866, 2974), 'langchain.text_splitter.RecursiveCharacterTextSplitter.from_tiktoken_encoder', 'RecursiveCharacterTextSplitter.from_tiktoken_encoder', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size,\n    chunk_overlap=chunk_overlap)\n', (2918, 2974), False, 'from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter\n'), ((3091, 3190), 'langchain.text_splitter.CharacterTextSplitter.from_tiktoken_encoder', 'CharacterTextSplitter.from_tiktoken_encoder', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size,\n    chunk_overlap=chunk_overlap)\n', (3134, 3190), False, 'from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter\n'), ((3303, 3374), 'langchain.text_splitter.TokenTextSplitter', 'LCTokenTextSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n', (3322, 3374), True, 'from langchain.text_splitter import TokenTextSplitter as LCTokenTextSplitter\n')]