Spaces:

umarbalak
/

QueryMind

Sleeping

update patch

84579b9 over 1 year ago

9.22 kB

# First import our patch before any gradio imports
try:
    from patch_gradio import patched_get_type
except ImportError:
    pass

# Now import the rest of the packages
import gradio as gr
from pathlib import Path
from tempfile import mkdtemp

# LangChain & Embedding/LLM
from langchain_community.vectorstores import FAISS
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_huggingface import HuggingFaceEndpoint
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import PromptTemplate
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Custom Data Extractor
from data_extraction import DataExtractor

# Constants
EMBED_MODEL_ID = "BAAI/bge-large-en-v1.5"
GEN_MODEL_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1"
TOP_K = 5
TEMP_DIR = Path(mkdtemp())
milvus_uri = str(TEMP_DIR / "docling.db")

# Initialize
extractor = DataExtractor()
embedding = HuggingFaceEmbeddings(model_name=EMBED_MODEL_ID)
llm = HuggingFaceEndpoint(repo_id=GEN_MODEL_ID, task="text-generation")

template = """You are a helpful assistant. Based on the following context, answer the user's query.
If you don't know the answer, just say that you don't know. Don't try to make up an answer.
Use the given context strictly to answer the question.
<context>
{context}
</context>
Question: {input}
Helpful Answer:"""

rag_prompt = PromptTemplate.from_template(template)

# App State
vectorstore = None
retriever = None  # Store the retriever separately
rag_chain = None

# Function to process using direct extraction methods
def process_direct(file_path):
    text = ""
    if file_path.suffix.lower() == ".pdf":
        text = extractor.extract_text_from_pdf(file_path)
    elif file_path.suffix.lower() == ".docx":
        text = extractor.extract_text_from_doc(file_path)
    elif file_path.suffix.lower() == ".txt":
        text = extractor.extract_text_from_txt(file_path)
    else:
        return None, "Unsupported file format. Please use PDF, DOCX, or TXT."
    
    # Clean text
    text = extractor.clean_text(text)
    
    # Split text using RecursiveCharacterTextSplitter
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len
    )
    
    from langchain_core.documents import Document
    doc = Document(
        page_content=text, 
        metadata={
            "source": str(file_path),
            "filename": file_path.name
        }
    )
    return text_splitter.split_documents([doc]), "Direct extraction completed successfully."

# Function to process using docling
def process_docling(file_path):
    try:
        if file_path.suffix.lower() == ".pdf":
            # Convert PDF to markdown
            markdown_path = TEMP_DIR / "temp.md"
            extractor.pdf_to_markdown(file_path, markdown_path)
            # Load markdown
            markdown_raw = extractor.load_markdown_file(markdown_path)
        elif file_path.suffix.lower() == ".md":
            markdown_raw = extractor.load_markdown_file(file_path)
        else:
            return None, "Only PDF or Markdown files are supported for docling extraction."
        
        # Update metadata with filename
        for doc in markdown_raw:
            doc.metadata.update({
                "filename": file_path.name
            })
            
        # Split markdown
        splits = extractor.spit_markdown(markdown_raw)
        
        # Update metadata for all splits
        for split in splits:
            split.metadata.update({
                "filename": file_path.name
            })
            
        return splits, "Docling extraction completed successfully."
    except Exception as e:
        return None, f"Error during docling extraction: {str(e)}"

# Step 1: File Upload & Index
def process_file(file, extraction_method):
    global vectorstore, retriever, rag_chain
    
    if file is None:
        return "Please upload a file first."

    file_path = Path(file.name)
    
    # Process based on selected extraction method
    if extraction_method == "Direct Extraction":
        data_splits, message = process_direct(file_path)
    else:  # Docling Extraction
        data_splits, message = process_docling(file_path)
    
    if data_splits is None:
        return message
    
    # Create vector store
    vectorstore = FAISS.from_documents(documents=data_splits, embedding=embedding)
    
    # Create retriever and chain
    retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})
    question_answer_chain = create_stuff_documents_chain(llm, prompt=rag_prompt)
    rag_chain = create_retrieval_chain(retriever, question_answer_chain)

    # Count chunks created
    chunk_count = len(data_splits)
    
    return f"File processed using {extraction_method}. {message} Created {chunk_count} document chunks. You can now ask questions."

# Format source information
def format_sources(docs):
    source_info = []
    seen_sources = set()
    
    for i, doc in enumerate(docs, 1):
        # Extract source information
        source = doc.metadata.get("source", "Unknown")
        filename = doc.metadata.get("filename", Path(source).name if source != "Unknown" else "Unknown")
        
        # Get header information if available (from docling extraction)
        header_info = []
        for level in range(1, 4):
            header = doc.metadata.get(f"Header_{level}", "")
            if header:
                header_info.append(header)
        
        # Create a unique identifier for this source
        source_id = f"{filename}{'_' + '_'.join(header_info) if header_info else ''}"
        
        if source_id not in seen_sources:
            seen_sources.add(source_id)
            
            # Format the source information
            source_text = f"**Source {i}: {filename}**"
            if header_info:
                source_text += f"\n - Section: {' > '.join(header_info)}"
            
            # Add the full text content without truncation
            source_text += f"\n - Content: \"{doc.page_content}\""
            
            source_info.append(source_text)
    
    return "\n\n".join(source_info)

# Step 2: Chat with RAG
def chat_with_doc(query):
    if not rag_chain:
        return "Upload and process a document first.", ""
    
    # Use the stored retriever directly to get documents
    if not retriever:
        return "Retriever not initialized. Please process a document first.", ""
    
    # Get documents directly from retriever
    retrieved_docs = retriever.get_relevant_documents(query)
    
    # Then invoke the chain with the query
    result = rag_chain.invoke({"input": query})
    answer = result.get("answer", "No answer found.")
    
    # Format source information
    source_text = format_sources(retrieved_docs)
    
    return answer, source_text

# Graceful error handling wrapper for Gradio functions
def safe_process_file(file, extraction_method):
    try:
        return process_file(file, extraction_method)
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"Error processing file: {str(e)}\n{error_trace}")
        return f"Error processing file: {str(e)}"

def safe_chat_with_doc(query):
    try:
        return chat_with_doc(query)
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"Error processing query: {str(e)}\n{error_trace}")
        return f"Error processing query: {str(e)}", ""

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# Document Q&A System")
    
    with gr.Tab("Upload Document"):
        file_input = gr.File(label="Upload File", file_types=[".pdf", ".md", ".docx", ".txt"])
        extraction_method = gr.Radio(
            ["Direct Extraction", "Docling Extraction"], 
            label="Extraction Method",
            value="Direct Extraction",
            info="Direct extraction is faster but simpler. Docling extraction preserves document structure better."
        )
        process_btn = gr.Button("Process Document")
        file_output = gr.Textbox(label="Processing Status", interactive=False)
        
        process_btn.click(
            fn=safe_process_file,  # Use the wrapped safe function
            inputs=[file_input, extraction_method],
            outputs=file_output
        )
    
    with gr.Tab("Chat with Document"):
        chat_input = gr.Textbox(label="Ask a Question", placeholder="Type your question...")
        ask_btn = gr.Button("Ask")
        chat_output = gr.Textbox(label="Answer", interactive=False)
        sources_output = gr.Textbox(label="Sources", interactive=False)
        
        ask_btn.click(
            fn=safe_chat_with_doc,  # Use the wrapped safe function
            inputs=chat_input,
            outputs=[chat_output, sources_output]
        )
        
        chat_input.submit(
            fn=safe_chat_with_doc,  # Use the wrapped safe function
            inputs=chat_input,
            outputs=[chat_output, sources_output]
        )

if __name__ == "__main__":
    app.launch(debug=True)

	# First import our patch before any gradio imports
	try:
	from patch_gradio import patched_get_type
	except ImportError:
	pass

	# Now import the rest of the packages
	import gradio as gr
	from pathlib import Path
	from tempfile import mkdtemp

	# LangChain & Embedding/LLM
	from langchain_community.vectorstores import FAISS
	from langchain_huggingface.embeddings import HuggingFaceEmbeddings
	from langchain_huggingface import HuggingFaceEndpoint
	from langchain.chains import create_retrieval_chain
	from langchain.chains.combine_documents import create_stuff_documents_chain
	from langchain_core.prompts import PromptTemplate
	from langchain_text_splitters import RecursiveCharacterTextSplitter

	# Custom Data Extractor
	from data_extraction import DataExtractor

	# Constants
	EMBED_MODEL_ID = "BAAI/bge-large-en-v1.5"
	GEN_MODEL_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1"
	TOP_K = 5
	TEMP_DIR = Path(mkdtemp())
	milvus_uri = str(TEMP_DIR / "docling.db")

	# Initialize
	extractor = DataExtractor()
	embedding = HuggingFaceEmbeddings(model_name=EMBED_MODEL_ID)
	llm = HuggingFaceEndpoint(repo_id=GEN_MODEL_ID, task="text-generation")

	template = """You are a helpful assistant. Based on the following context, answer the user's query.
	If you don't know the answer, just say that you don't know. Don't try to make up an answer.
	Use the given context strictly to answer the question.

	<context>
	{context}
	</context>

	Question: {input}
	Helpful Answer:"""

	rag_prompt = PromptTemplate.from_template(template)

	# App State
	vectorstore = None
	retriever = None # Store the retriever separately
	rag_chain = None

	# Function to process using direct extraction methods
	def process_direct(file_path):
	text = ""
	if file_path.suffix.lower() == ".pdf":
	text = extractor.extract_text_from_pdf(file_path)
	elif file_path.suffix.lower() == ".docx":
	text = extractor.extract_text_from_doc(file_path)
	elif file_path.suffix.lower() == ".txt":
	text = extractor.extract_text_from_txt(file_path)
	else:
	return None, "Unsupported file format. Please use PDF, DOCX, or TXT."

	# Clean text
	text = extractor.clean_text(text)

	# Split text using RecursiveCharacterTextSplitter
	text_splitter = RecursiveCharacterTextSplitter(
	chunk_size=1000,
	chunk_overlap=200,
	length_function=len
	)

	from langchain_core.documents import Document
	doc = Document(
	page_content=text,
	metadata={
	"source": str(file_path),
	"filename": file_path.name
	}
	)
	return text_splitter.split_documents([doc]), "Direct extraction completed successfully."

	# Function to process using docling
	def process_docling(file_path):
	try:
	if file_path.suffix.lower() == ".pdf":
	# Convert PDF to markdown
	markdown_path = TEMP_DIR / "temp.md"
	extractor.pdf_to_markdown(file_path, markdown_path)
	# Load markdown
	markdown_raw = extractor.load_markdown_file(markdown_path)
	elif file_path.suffix.lower() == ".md":
	markdown_raw = extractor.load_markdown_file(file_path)
	else:
	return None, "Only PDF or Markdown files are supported for docling extraction."

	# Update metadata with filename
	for doc in markdown_raw:
	doc.metadata.update({
	"filename": file_path.name
	})

	# Split markdown
	splits = extractor.spit_markdown(markdown_raw)

	# Update metadata for all splits
	for split in splits:
	split.metadata.update({
	"filename": file_path.name
	})

	return splits, "Docling extraction completed successfully."
	except Exception as e:
	return None, f"Error during docling extraction: {str(e)}"

	# Step 1: File Upload & Index
	def process_file(file, extraction_method):
	global vectorstore, retriever, rag_chain

	if file is None:
	return "Please upload a file first."

	file_path = Path(file.name)

	# Process based on selected extraction method
	if extraction_method == "Direct Extraction":
	data_splits, message = process_direct(file_path)
	else: # Docling Extraction
	data_splits, message = process_docling(file_path)

	if data_splits is None:
	return message

	# Create vector store
	vectorstore = FAISS.from_documents(documents=data_splits, embedding=embedding)

	# Create retriever and chain
	retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})
	question_answer_chain = create_stuff_documents_chain(llm, prompt=rag_prompt)
	rag_chain = create_retrieval_chain(retriever, question_answer_chain)

	# Count chunks created
	chunk_count = len(data_splits)

	return f"File processed using {extraction_method}. {message} Created {chunk_count} document chunks. You can now ask questions."

	# Format source information
	def format_sources(docs):
	source_info = []
	seen_sources = set()

	for i, doc in enumerate(docs, 1):
	# Extract source information
	source = doc.metadata.get("source", "Unknown")
	filename = doc.metadata.get("filename", Path(source).name if source != "Unknown" else "Unknown")

	# Get header information if available (from docling extraction)
	header_info = []
	for level in range(1, 4):
	header = doc.metadata.get(f"Header_{level}", "")
	if header:
	header_info.append(header)

	# Create a unique identifier for this source
	source_id = f"{filename}{'_' + '_'.join(header_info) if header_info else ''}"

	if source_id not in seen_sources:
	seen_sources.add(source_id)

	# Format the source information
	source_text = f"Source {i}: {filename}"
	if header_info:
	source_text += f"\n - Section: {' > '.join(header_info)}"

	# Add the full text content without truncation
	source_text += f"\n - Content: \"{doc.page_content}\""

	source_info.append(source_text)

	return "\n\n".join(source_info)

	# Step 2: Chat with RAG
	def chat_with_doc(query):
	if not rag_chain:
	return "Upload and process a document first.", ""

	# Use the stored retriever directly to get documents
	if not retriever:
	return "Retriever not initialized. Please process a document first.", ""

	# Get documents directly from retriever
	retrieved_docs = retriever.get_relevant_documents(query)

	# Then invoke the chain with the query
	result = rag_chain.invoke({"input": query})
	answer = result.get("answer", "No answer found.")

	# Format source information
	source_text = format_sources(retrieved_docs)

	return answer, source_text

	# Graceful error handling wrapper for Gradio functions
	def safe_process_file(file, extraction_method):
	try:
	return process_file(file, extraction_method)
	except Exception as e:
	import traceback
	error_trace = traceback.format_exc()
	print(f"Error processing file: {str(e)}\n{error_trace}")
	return f"Error processing file: {str(e)}"

	def safe_chat_with_doc(query):
	try:
	return chat_with_doc(query)
	except Exception as e:
	import traceback
	error_trace = traceback.format_exc()
	print(f"Error processing query: {str(e)}\n{error_trace}")
	return f"Error processing query: {str(e)}", ""

	# Gradio Interface
	with gr.Blocks() as app:
	gr.Markdown("# Document Q&A System")

	with gr.Tab("Upload Document"):
	file_input = gr.File(label="Upload File", file_types=[".pdf", ".md", ".docx", ".txt"])
	extraction_method = gr.Radio(
	["Direct Extraction", "Docling Extraction"],
	label="Extraction Method",
	value="Direct Extraction",
	info="Direct extraction is faster but simpler. Docling extraction preserves document structure better."
	)
	process_btn = gr.Button("Process Document")
	file_output = gr.Textbox(label="Processing Status", interactive=False)

	process_btn.click(
	fn=safe_process_file, # Use the wrapped safe function
	inputs=[file_input, extraction_method],
	outputs=file_output
	)

	with gr.Tab("Chat with Document"):
	chat_input = gr.Textbox(label="Ask a Question", placeholder="Type your question...")
	ask_btn = gr.Button("Ask")
	chat_output = gr.Textbox(label="Answer", interactive=False)
	sources_output = gr.Textbox(label="Sources", interactive=False)

	ask_btn.click(
	fn=safe_chat_with_doc, # Use the wrapped safe function
	inputs=chat_input,
	outputs=[chat_output, sources_output]
	)

	chat_input.submit(
	fn=safe_chat_with_doc, # Use the wrapped safe function
	inputs=chat_input,
	outputs=[chat_output, sources_output]
	)

	if __name__ == "__main__":
	app.launch(debug=True)