# 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.
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, config, 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()
