# TODO: return all pages used to form answer # TODO: question samples # TEST: with and without GPU instance # TODO: visual questions on page image (in same app)? # expose more parameters import torch from llama_index.llms.huggingface import HuggingFaceLLM from llama_index.embeddings.huggingface import HuggingFaceEmbedding from llama_index.core import SimpleDirectoryReader from llama_index.core import VectorStoreIndex, SummaryIndex from llama_index.core.prompts import PromptTemplate from llama_index.core import Settings from PIL import Image import gradio as gr CHEAPMODE = torch.cuda.is_available() # LLM = "HuggingFaceH4/zephyr-7b-alpha" if not CHEAPMODE else "microsoft/phi-2" config = { # "LLM": "meta-llama/Meta-Llama-3-8B", # "LLM": "google/gemma-2b", # "LLM": "microsoft/phi-2", "LLM": "HuggingFaceH4/zephyr-7b-alpha", "embeddings": "BAAI/bge-small-en-v1.5", "similarity_top_k": 2, "context_window": 2048, "max_new_tokens": 200, "temperature": 0.7, "top_k": 5, "top_p": 0.95, "chunk_size": 512, "chunk_overlap": 50, } def center_element(el): return f"
{el}
" title = "Ask my thesis: Intelligent Automation for AI-Driven Document Understanding" title = center_element(title) description = """Chat with the thesis manuscript by asking questions and receive answers with reference to the page.
Ideogram image generated with prompt engineering
Click the visual above to be redirected to the PDF of the manuscript. Technology used: [Llama-index](https://www.llamaindex.ai/), OS LLMs from HuggingFace Spoiler: a quickly hacked together RAG application with a >1B LLM and online vector store can be quite slow on a 290 page document ⏳ (10s+) """ description = center_element(description) def messages_to_prompt(messages): prompt = "" for message in messages: if message.role == "system": m = "You are an expert in the research field of document understanding, bayesian deep learning and neural networks." prompt += f"<|system|>\n{m}\n" elif message.role == "user": prompt += f"<|user|>\n{message.content}\n" elif message.role == "assistant": prompt += f"<|assistant|>\n{message.content}\n" # ensure we start with a system prompt, insert blank if needed if not prompt.startswith("<|system|>\n"): prompt = "<|system|>\n\n" + prompt # add final assistant prompt prompt = prompt + "<|assistant|>\n" return prompt def load_RAG_pipeline(config): # LLM quantization_config = None # dirty fix for CPU/GPU support if torch.cuda.is_available(): from transformers import BitsAndBytesConfig quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, ) llm = HuggingFaceLLM( model_name=config["LLM"], tokenizer_name=config["LLM"], query_wrapper_prompt=PromptTemplate("<|system|>\n\n<|user|>\n{query_str}\n<|assistant|>\n"), context_window=config["context_window"], max_new_tokens=config["max_new_tokens"], model_kwargs={"quantization_config": quantization_config}, # tokenizer_kwargs={}, generate_kwargs={"temperature": config["temperature"], "top_k": config["top_k"], "top_p": config["top_p"]}, messages_to_prompt=messages_to_prompt, device_map="auto", ) # Llama-index Settings.llm = llm Settings.embed_model = HuggingFaceEmbedding(model_name=config["embeddings"]) print(Settings) Settings.chunk_size = config["chunk_size"] Settings.chunk_overlap = config["chunk_overlap"] # raw data documents = SimpleDirectoryReader("assets/txts").load_data() vector_index = VectorStoreIndex.from_documents(documents) # summary_index = SummaryIndex.from_documents(documents) # vector_index.persist(persist_dir="vectors") # https://docs.llamaindex.ai/en/v0.10.17/understanding/storing/storing.html query_engine = vector_index.as_query_engine(response_mode="compact", similarity_top_k=config["similarity_top_k"]) return query_engine default_query_engine = load_RAG_pipeline(config) # These are placeholder functions to simulate the behavior of the RAG setup. # You would need to implement these with the actual logic to retrieve and generate answers based on the document. def get_answer(question, query_engine=default_query_engine): # Here you should implement the logic to generate an answer based on the question and the document. # For example, you could use a machine learning model for RAG. # answer = "This is a placeholder answer." # https://docs.llamaindex.ai/en/stable/module_guides/supporting_modules/settings/#setting-local-configurations response = query_engine.query(question) print(f"A: {response}") return response def get_answer_page(response): # Implement logic to retrieve the page number or an image of the page with the answer. # best image best_match = response.source_nodes[0].metadata["file_path"] answer_page = int(best_match[-8:-4]) image = Image.open(best_match.replace("txt", "png")) return image, f"Navigate to page {answer_page}" # Create the gr.Interface function def ask_my_thesis( question, similarity_top_k=config["similarity_top_k"], context_window=config["context_window"], max_new_tokens=config["max_new_tokens"], temperature=config["temperature"], top_k=config["top_k"], top_p=config["top_p"], chunk_size=config["chunk_size"], chunk_overlap=config["chunk_overlap"], ): print(f"Got Q: {question}") query_engine = default_query_engine # if any change in kwargs # Check if any of the kwargs have changed if ( temperature != config["temperature"] or top_p != config["top_p"] or max_new_tokens != config["max_new_tokens"] # or LLM != config["LLM"] # or embeddings != config["embeddings"] or similarity_top_k != config["similarity_top_k"] or context_window != config["context_window"] or top_k != config["top_k"] or chunk_size != config["chunk_size"] or chunk_overlap != config["chunk_overlap"] ): # Update the config dictionary with the new values config["temperature"] = temperature config["top_p"] = top_p config["max_new_tokens"] = max_new_tokens # config["LLM"] = LLM # config["embeddings"] = embeddings config["similarity_top_k"] = similarity_top_k config["context_window"] = context_window config["top_k"] = top_k config["chunk_size"] = chunk_size config["chunk_overlap"] = chunk_overlap query_engine = load_RAG_pipeline(config) answer = get_answer(question, query_engine=query_engine) image, answer_page = get_answer_page(answer) return answer.response, image, answer_page # Set up the interface options based on the design in the image. output_image = gr.Image(label="Answer Page") # examples examples = [ ["What model is state-of-the-art on DUDE?"], ["Why is knowledge distillation interesting?"], ["What is ANLS?"], ] # Define additional Gradio input elements additional_inputs = [ # gr.Input("text", label="Question"), # gr.Input("text", label="LLM", value=config["LLM"]), # gr.Input("text", label="Embeddings", value=config["embeddings"]), gr.Slider(1, 5, value=config["similarity_top_k"], label="Similarity Top K", step=1), gr.Slider(512, 8048, value=config["context_window"], label="Context Window"), gr.Slider(20, 500, value=config["max_new_tokens"], label="Max New Tokens"), gr.Slider(0, 1, value=config["temperature"], label="Temperature"), gr.Slider(1, 10, value=config["top_k"], label="Top K", step=1), gr.Slider(0, 1, value=config["top_p"], label="Nucleus Sampling"), gr.Slider(128, 4024, value=config["chunk_size"], label="Chunk Size"), gr.Slider(0, 200, value=config["chunk_overlap"], label="Chunk Overlap"), ] iface = gr.Interface( fn=ask_my_thesis, inputs=[gr.Textbox(label="Question", placeholder="Type your question here...")], additional_inputs=additional_inputs, outputs=[gr.Textbox(label="Answer"), output_image, gr.Label()], examples=examples, title=title, description=description, allow_flagging="auto", cache_examples=True, ) # https://github.com/gradio-app/gradio/issues/4309 # https://discuss.huggingface.co/t/add-background-image/16381/4 background image # Start the application. if __name__ == "__main__": iface.launch()