app.py

import tempfile
import os
import tiktoken
import streamlit as st

from llama_index.core import (
    VectorStoreIndex,
    Settings,
)

from llama_parse import LlamaParse
from streamlit_pdf_viewer import pdf_viewer

class MistralTokens:
    """
    Returns tokens for MistralAI models.
    
    See: https://docs.mistral.ai/guides/tokenization/
    """
    def __init__(self, llm_name):
        from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
        if 'open-mistral-nemo' in llm_name:
            self.tokenizer = MistralTokenizer.v3(is_tekken=True)
        else:
            # This might work for all models, but their documentation is unclear.
            self.tokenizer = MistralTokenizer.from_model(llm_name)

    def __call__(self, input):
        """This returns all the tokens indices in a list since LlamaIndex seems to count by calling `len()` on the tokenizer function."""
        from mistral_common.protocol.instruct.messages import UserMessage
        from mistral_common.protocol.instruct.request import ChatCompletionRequest

        return self.tokenizer.encode_chat_completion(
            ChatCompletionRequest(
                tools=[], 
                messages=[
                    UserMessage(content=input)
                ]
            )
        ).tokens

class GeminiTokens:
    """
    Returns tokens for Gemini models.
    
    See: https://medium.com/google-cloud/counting-gemini-text-tokens-locally-with-the-vertex-ai-sdk-78979fea6244
    """
    def __init__(self, llm_name):
        from vertexai.preview import tokenization
        self.tokenizer = tokenization.get_tokenizer_for_model(llm_name)

    def __call__(self, input):
        """This returns all the tokens in a list since LlamaIndex seems to count by calling `len()` on the tokenizer function."""
        tokens = []
        for list in self.tokenizer.compute_tokens(input).token_info_list:
            tokens += list.tokens
        return tokens
        
def main():
    with st.sidebar:
        st.title('Document Summarization and QA System')

        # Select Provider
        provider = st.selectbox(
            label="Select LLM Provider",
            options=['google', 'huggingface', 'mistralai', 'openai'],
            index=3
        )

        # Select LLM
        if provider == 'google':
            llm_list = ['gemini-1.0-pro', 'gemini-1.5-flash', 'gemini-1.5-pro']
        elif provider == 'huggingface':
            llm_list = []
        elif provider == 'mistralai':
            llm_list = ["mistral-large-latest", "open-mistral-nemo-latest"]
        elif provider == 'openai':
            llm_list = ['gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'gpt-4o', 'gpt-4o-mini']
        else:
            llm_list = []

        if provider == 'huggingface':
            llm_name = st.text_input(
                "Enter LLM namespace/model-name",
                value="HuggingFaceH4/zephyr-7b-alpha",
            )

            # Also give the user the option for different embedding models, too
            embed_name = st.text_input(
                label="Enter embedding namespace/model-name",
                value="BAAI/bge-small-en-v1.5",
            )
        else:
            llm_name = st.selectbox(
                label="Select LLM Model",
                options=llm_list,
                index=0
            )

        # Temperature
        temperature = st.slider(
            "Temperature",
            min_value=0.0, 
            max_value=1.0, 
            value=0.0, 
            step=0.05, 
        )

        # Enter Parsing API Key
        parse_key = st.text_input(
            "Enter your LlamaParse API Key",
            value=None
        )

        # Enter LLM API Key
        llm_key = st.text_input(
            "Enter your LLM provider API Key",
            value=None,
        )

        # Create LLM
        # Global tokenization needs to be consistent with LLM for token counting
        # https://docs.llamaindex.ai/en/stable/module_guides/models/llms/
        if llm_key is not None:
            if provider == 'google':
                from llama_index.llms.gemini import Gemini
                from llama_index.embeddings.gemini import GeminiEmbedding
                max_output_tokens = 8192 # https://firebase.google.com/docs/vertex-ai/gemini-models
                
                os.environ['GOOGLE_API_KEY'] = str(llm_key)
                Settings.llm = Gemini(
                    model=f"models/{llm_name}",
                    token=os.environ.get("GOOGLE_API_KEY"),
                    temperature=temperature,
                    max_tokens=max_output_tokens
                )
                Settings.tokenizer = GeminiTokens(llm_name) 
                Settings.num_output = max_output_tokens
                Settings.embed_model = GeminiEmbedding(
                    model_name="models/text-embedding-004", api_key=os.environ.get("GOOGLE_API_KEY") #, title="this is a document"
                )
                if llm_name == 'gemini-1.0-pro':
                    total_token_limit = 32760
                else:
                    total_token_limit = 1e6
                Settings.context_window = total_token_limit -  max_output_tokens # Gemini counts total tokens
            elif provider == 'huggingface':
                if llm_name is not None and embed_name is not None:
                    from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI 
                    from llama_index.embeddings.huggingface import HuggingFaceInferenceAPIEmbedding
                    from transformers import AutoTokenizer

                    max_output_tokens = 2048 # Just a generic value

                    os.environ['HFTOKEN'] = str(llm_key)
                    Settings.llm = HuggingFaceInferenceAPI(
                        model_name=llm_name, 
                        token=os.environ.get("HFTOKEN"),
                        temperature=temperature,
                        max_tokens=max_output_tokens
                    )
                    Settings.tokenizer = AutoTokenizer.from_pretrained(
                        llm_name,
                        token=os.environ.get("HFTOKEN"),
                    )
                    Settings.num_output = max_output_tokens
                    Settings.embed_model = HuggingFaceInferenceAPIEmbedding(
                        model_name=embed_name
                    )
                    Settings.context_window = 4096 # Just a generic value
            elif provider == 'mistralai':
                from llama_index.llms.mistralai import MistralAI
                from llama_index.embeddings.mistralai import MistralAIEmbedding
                max_output_tokens = 8192 # Based on internet consensus since this is not well documented
                
                os.environ['MISTRAL_API_KEY'] = str(llm_key)
                Settings.llm = MistralAI(
                    model=llm_name, 
                    temperature=temperature,
                    max_tokens=max_output_tokens,
                    random_seed=42,
                    safe_mode=True
                )
                Settings.tokenizer = MistralTokens(llm_name)
                Settings.num_output = max_output_tokens
                Settings.embed_model = MistralAIEmbedding(
                    model_name="mistral-embed", 
                    api_key=os.environ.get("MISTRAL_API_KEY")
                )
                Settings.context_window = 128000 # 128k for flagship models - doesn't seem to count input tokens
            elif provider == 'openai':
                from llama_index.llms.openai import OpenAI
                from llama_index.embeddings.openai import OpenAIEmbedding

                # https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
                if llm_name == 'gpt-3.5-turbo':
                    max_output_tokens = 4096
                    context_window = 16385
                elif llm_name == 'gpt-4' :
                    max_output_tokens = 8192
                    context_window = 8192
                elif llm_name == 'gpt-4-turbo'
                    max_output_tokens = 4096
                    context_window = 128000
                elif llm_name == 'gpt-4o':
                    max_output_tokens = 4096
                    context_window = 128000
                elif llm_name == 'gpt-4o-mini':
                    max_output_tokens = 16384
                    context_window = 128000

                os.environ["OPENAI_API_KEY"] = str(llm_key)
                Settings.llm = OpenAI(
                    model=llm_name, 
                    temperature=temperature,
                    max_tokens=max_output_tokens
                )
                Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
                Settings.num_output = max_output_tokens
                Settings.embed_model = OpenAIEmbedding()
                Settings.context_window = context_window 
            else:
                raise NotImplementedError(f"{provider} is not supported yet")

        uploaded_file = st.file_uploader(
            "Choose a PDF file to upload", 
            type=['pdf'], 
            accept_multiple_files=False
        )

        parsed_document = None
        if uploaded_file is not None:
            # Parse the file
            parser = LlamaParse(
                api_key=parse_key,  # Can also be set in your env as LLAMA_CLOUD_API_KEY
                result_type="text"  # "markdown" and "text" are available
            )

            # Create a temporary directory to save the file then load and parse it
            temp_dir = tempfile.TemporaryDirectory()
            temp_filename = os.path.join(temp_dir.name, uploaded_file.name)
            with open(temp_filename, "wb") as f:
                f.write(uploaded_file.getvalue())
            parsed_document = parser.load_data(temp_filename)
            temp_dir.cleanup()

    col1, col2 = st.columns(2)

    with col2:
        tab1, tab2 = st.tabs(["Uploaded File", "Parsed File",])

        with tab1:
            if uploaded_file is not None: # Display the pdf
                bytes_data = uploaded_file.getvalue()
                pdf_viewer(input=bytes_data, width=700)    
        
        with tab2:
            if parsed_document is not None: # Showed the raw parsing result
                st.write(parsed_document)

    with col1:
        st.markdown(
            """
            # Instructions

            1. Obtain an [API Key](https://cloud.llamaindex.ai/api-key) from LlamaParse to parse your document. 
            2. Obtain a similar API Key from your preferred LLM provider. Note, if you are using [Hugging Face](https://huggingface.co/models) you may need to request access to a model if it is gated.
            3. Make selections at the left and upload a document to use as context.
            4. Begin asking questions below!
            """
        )

        st.divider()

        prompt_txt = 'You are a trusted scientific expert that only responds truthfully to inquiries. Summarize this document in a 3-5 sentences.'
        prompt = st.text_area(
            label="Enter your query.",
            key="prompt_widget",
            value=prompt_txt
        )

        run = st.button("Answer", type="primary")

        if parsed_document is not None and run:
            index = VectorStoreIndex.from_documents(parsed_document)
            query_engine = index.as_query_engine()
            response = query_engine.query(prompt)
            st.write(response.response)

if __name__ == '__main__':
    # Global configurations
    # from llama_index.core import set_global_handler
    # set_global_handler("langfuse")
    # Also add API Key for this if using

    st.set_page_config(layout="wide")

    main()