Spaces:

KalooAI
/

PDF-Analyzer

Running

Update app.py

451bc89 verified 9 months ago

11.7 kB

import os
import time
import gradio as gr
import numpy as np
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
from huggingface_hub import InferenceClient
from huggingface_hub.errors import HfHubHTTPError

# ============================================================
# CONFIG
# ============================================================

HF_TOKEN = os.getenv("HF_TOKEN", None)

# Free model compatible with new HF API
#LLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
#LLM_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
# LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.3" Does not work
# LLM_MODEL ="HuggingFaceH4/zephyr-7b-beta"
# LLM_MODEL = "NousResearch/Hermes-2-Pro-Mistral-7B" Does not work

# Were Great LLM
# LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
# LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"

# LLM model
# LLM_MODEL = "deepseek-ai/DeepSeek-V3"
LLM_MODEL = 'Qwen/Qwen2.5-14B-Instruct'

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
embedding_model = SentenceTransformer(EMBEDDING_MODEL)

# IMPORTANT: do NOT pass model=... here
llm = InferenceClient(token=HF_TOKEN)


# ============================================================
# SAFE LLM CALL (prevents 429 errors)
# ============================================================

def safe_chat_completion(llm, **kwargs):
    for attempt in range(5):
        try:
            return llm.chat.completions.create(**kwargs)
        except HfHubHTTPError as e:
            if "429" in str(e):
                sleep_time = 2 ** attempt
                print(f"[Rate limit] Retrying in {sleep_time}s...")
                time.sleep(sleep_time)
            else:
                raise
    raise RuntimeError("Too many retries after rate limits.")


# ============================================================
# PDF PROCESSING
# ============================================================

def read_pdf(file_obj):
    reader = PdfReader(file_obj)
    pages = []
    for page in reader.pages:
        text = page.extract_text()
        if text:
            pages.append(text)
    return "\n".join(pages)


def chunk_text(text, chunk_size=800, overlap=200):
    words = text.split()
    chunks = []
    start = 0

    while start < len(words):
        end = start + chunk_size
        chunk = " ".join(words[start:end])
        chunks.append(chunk)
        start += chunk_size - overlap

    return chunks


# ============================================================
# PER-FILE RETRIEVAL
# ============================================================

def retrieve_context_per_file(question, state, k=10):
    """
    For each file, retrieve top-k chunks separately,
    then merge all selected chunks into a single context.
    """
    if not state["files"]:
        return ""

    q_emb = embedding_model.encode([question], convert_to_numpy=True)[0]
    all_selected = []

    for filename, data in state["files"].items():
        chunks = data["chunks"]
        embeddings = data["embeddings"]

        if len(chunks) == 0:
            continue

        # Compute L2 distances to the query embedding
        distances = np.linalg.norm(embeddings - q_emb, axis=1)

        # Get top-k indices for this file
        top_k_idx = distances.argsort()[: min(k, len(chunks))]

        # Tag chunks with filename for transparency
        selected = [f"[{filename}]\n{chunks[i]}" for i in top_k_idx]
        all_selected.extend(selected)

    if not all_selected:
        return ""

    return "\n\n---\n\n".join(all_selected)


# ============================================================
# LLM ANSWERING
# ============================================================

def generate_answer(question, context):
    if context:
        prompt = (
            "You are a helpful assistant answering questions strictly from the provided context.\n"
            "If the answer is not in the context, say you don't know.\n\n"
            f"Context:\n{context}\n\n"
            f"Question: {question}\nAnswer:"
        )
    else:
        prompt = (
            "You are a helpful assistant. The user asked a question but no document context is available.\n\n"
            f"Question: {question}\nAnswer:"
        )

    response = safe_chat_completion(
        llm,
        model=LLM_MODEL,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1024,
        temperature=0.25,
        top_p=0.9,
        frequency_penalty=0.1,
        presence_penalty=0.0,
    )

    return response.choices[0].message["content"].strip()


# ============================================================
# GRADIO CALLBACKS
# ============================================================

def load_pdf(files, state):
    """
    files: current list from the UI (after add/remove).
    We rebuild state from scratch to stay in sync.
    state = {
        "files": {
            "filename.pdf": {
                "chunks": [...],
                "embeddings": np.array([...])
            },
            ...
        },
        "pdf_count": int,
        "total_chunks": int,
        "history": [...],
        "facts": [...]
    }
    """
    # If user cleared all files
    if not files:
        state["files"] = {}
        state["pdf_count"] = 0
        state["total_chunks"] = 0

        return (
            "No file uploaded.",
            state,
            gr.update(choices=[], value=None),
        )

    try:
        # Reset and rebuild from current files only
        state["files"] = {}
        state["pdf_count"] = 0
        state["total_chunks"] = 0

        # Deduplicate while preserving order
        unique_files = list(dict.fromkeys(files))

        for file in unique_files:
            filename = os.path.basename(str(file))

            text = read_pdf(file)
            chunks = chunk_text(text)

            if not chunks:
                continue

            embeddings = embedding_model.encode(chunks, convert_to_numpy=True)

            state["files"][filename] = {
                "chunks": chunks,
                "embeddings": embeddings,
            }

            state["pdf_count"] += 1
            state["total_chunks"] += len(chunks)

        status_msg = (
            f"{state['pdf_count']} PDF(s) loaded. "
            f"Total chunks: {state['total_chunks']}"
        )

        return (
            status_msg,
            state,
            gr.update(choices=list(state["files"].keys()), value=None),
        )

    except Exception as e:
        return (
            f"Error processing PDF: {e}",
            state,
            gr.update(choices=list(state["files"].keys()), value=None),
        )


def remove_pdf(filename, state):
    if filename not in state["files"]:
        return (
            f"{filename} not found.",
            state,
            gr.update(choices=list(state["files"].keys()), value=None),
        )

    # Remove the file entry
    state["files"].pop(filename, None)

    # Recompute counts
    pdf_count = len(state["files"])
    total_chunks = sum(len(data["chunks"]) for data in state["files"].values())

    state["pdf_count"] = pdf_count
    state["total_chunks"] = total_chunks

    status_msg = (
        f"Removed {filename}. "
        f"Total PDFs: {pdf_count}. Total chunks: {total_chunks}"
    )

    return (
        status_msg,
        state,
        gr.update(choices=list(state["files"].keys()), value=None),
    )


def answer_question(question, state):
    print("HF_TOKEN loaded:", HF_TOKEN is not None)

    if HF_TOKEN:
        print("HF_TOKEN length:", len(HF_TOKEN))
        print("HF_TOKEN preview:", HF_TOKEN[:6] + "..." + HF_TOKEN[-4:])
    else:
        print("HF_TOKEN is None or empty")

    question = question.strip()
    if not question:
        return "Please enter a question."

    if not state["files"]:
        return "Please upload at least one PDF first."

    # Simple persistent memory hook: "Remember that ..."
    lower_q = question.lower()
    if lower_q.startswith("remember that"):
        fact = question[len("remember that"):].strip(" .")
        if fact:
            state["facts"].append(fact)
            return f"Okay, I will remember that {fact}."
        else:
            return "What should I remember?"

    # RAG context (per-file balanced retrieval)
    rag_context = retrieve_context_per_file(question, state, k=10)

    # Conversation memory (last 5 turns)
    memory_context = ""
    if state["history"]:
        memory_context = "\n\n".join(
            [
                f"User: {h['user']}\nAssistant: {h['assistant']}"
                for h in state["history"][-5:]
            ]
        )

    # Persistent facts
    facts_context = ""
    if state["facts"]:
        facts_context = "User Facts:\n" + "\n".join(state["facts"]) + "\n\n"

    # Merge all contexts
    full_context_parts = []
    if facts_context:
        full_context_parts.append(facts_context)
    if memory_context:
        full_context_parts.append("Conversation Memory:\n" + memory_context)
    if rag_context:
        full_context_parts.append("Document Context:\n" + rag_context)

    full_context = "\n\n".join(full_context_parts)

    # Generate answer
    answer = generate_answer(question, full_context)

    # Save conversational memory
    state["history"].append({"user": question, "assistant": answer})

    return answer


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(title="PDF Q&A Chatbot") as demo:
    gr.Markdown("## 📄 PDF Q&A Chatbot\nUpload PDFs and ask questions about their content.")

    state = gr.State(
        {
            "files": {},          # filename -> {chunks, embeddings}
            "pdf_count": 0,
            "total_chunks": 0,
            "history": [],        # conversational memory
            "facts": [],          # persistent user facts
        }
    )

    with gr.Row():
        # LEFT COLUMN — PDF upload + removal
        with gr.Column(scale=1):
            pdf_input = gr.File(
                label="Upload PDF(s)",
                file_types=[".pdf"],
                file_count="multiple",  # proper multi-upload in Gradio 6
            )
            pdf_status = gr.Textbox(label="Status", interactive=False)

            file_selector = gr.Dropdown(
                label="Remove a PDF",
                choices=[],
                interactive=True,
            )
            remove_button = gr.Button("Remove Selected PDF")

        # RIGHT COLUMN — question + clear + answer
        with gr.Column(scale=2):
            with gr.Row():
                question = gr.Textbox(
                    label="Your question",
                    lines=2,
                    scale=3,
                )
                clear_button = gr.Button(
                    "Clear",
                    scale=1,
                    min_width=80,
                    elem_classes="clear-btn",
                )

            answer = gr.Textbox(label="Answer", lines=8)
            ask_button = gr.Button("Ask")

    # Bind events AFTER components exist
    pdf_input.change(
        load_pdf,
        inputs=[pdf_input, state],
        outputs=[pdf_status, state, file_selector],
    )

    remove_button.click(
        remove_pdf,
        inputs=[file_selector, state],
        outputs=[pdf_status, state, file_selector],
    )

    ask_button.click(
        answer_question,
        inputs=[question, state],
        outputs=answer,
    )



    clear_button.click(lambda: "", outputs=question)


if __name__ == "__main__":
    demo.launch(
        css=".clear-btn button { height: 100%; min-height: 52px; }"
    )


	import os
	import time
	import gradio as gr
	import numpy as np
	from PyPDF2 import PdfReader
	from sentence_transformers import SentenceTransformer
	from huggingface_hub import InferenceClient
	from huggingface_hub.errors import HfHubHTTPError

	# ============================================================
	# CONFIG
	# ============================================================

	HF_TOKEN = os.getenv("HF_TOKEN", None)

	# Free model compatible with new HF API
	#LLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
	#LLM_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
	# LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.3" Does not work
	# LLM_MODEL ="HuggingFaceH4/zephyr-7b-beta"
	# LLM_MODEL = "NousResearch/Hermes-2-Pro-Mistral-7B" Does not work

	# Were Great LLM
	# LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
	# LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"

	# LLM model
	# LLM_MODEL = "deepseek-ai/DeepSeek-V3"
	LLM_MODEL = 'Qwen/Qwen2.5-14B-Instruct'

	EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
	embedding_model = SentenceTransformer(EMBEDDING_MODEL)

	# IMPORTANT: do NOT pass model=... here
	llm = InferenceClient(token=HF_TOKEN)


	# ============================================================
	# SAFE LLM CALL (prevents 429 errors)
	# ============================================================

	def safe_chat_completion(llm, **kwargs):
	for attempt in range(5):
	try:
	return llm.chat.completions.create(**kwargs)
	except HfHubHTTPError as e:
	if "429" in str(e):
	sleep_time = 2 ** attempt
	print(f"[Rate limit] Retrying in {sleep_time}s...")
	time.sleep(sleep_time)
	else:
	raise
	raise RuntimeError("Too many retries after rate limits.")


	# ============================================================
	# PDF PROCESSING
	# ============================================================

	def read_pdf(file_obj):
	reader = PdfReader(file_obj)
	pages = []
	for page in reader.pages:
	text = page.extract_text()
	if text:
	pages.append(text)
	return "\n".join(pages)


	def chunk_text(text, chunk_size=800, overlap=200):
	words = text.split()
	chunks = []
	start = 0

	while start < len(words):
	end = start + chunk_size
	chunk = " ".join(words[start:end])
	chunks.append(chunk)
	start += chunk_size - overlap

	return chunks


	# ============================================================
	# PER-FILE RETRIEVAL
	# ============================================================

	def retrieve_context_per_file(question, state, k=10):
	"""
	For each file, retrieve top-k chunks separately,
	then merge all selected chunks into a single context.
	"""
	if not state["files"]:
	return ""

	q_emb = embedding_model.encode([question], convert_to_numpy=True)[0]
	all_selected = []

	for filename, data in state["files"].items():
	chunks = data["chunks"]
	embeddings = data["embeddings"]

	if len(chunks) == 0:
	continue

	# Compute L2 distances to the query embedding
	distances = np.linalg.norm(embeddings - q_emb, axis=1)

	# Get top-k indices for this file
	top_k_idx = distances.argsort()[: min(k, len(chunks))]

	# Tag chunks with filename for transparency
	selected = [f"[{filename}]\n{chunks[i]}" for i in top_k_idx]
	all_selected.extend(selected)

	if not all_selected:
	return ""

	return "\n\n---\n\n".join(all_selected)


	# ============================================================
	# LLM ANSWERING
	# ============================================================

	def generate_answer(question, context):
	if context:
	prompt = (
	"You are a helpful assistant answering questions strictly from the provided context.\n"
	"If the answer is not in the context, say you don't know.\n\n"
	f"Context:\n{context}\n\n"
	f"Question: {question}\nAnswer:"
	)
	else:
	prompt = (
	"You are a helpful assistant. The user asked a question but no document context is available.\n\n"
	f"Question: {question}\nAnswer:"
	)

	response = safe_chat_completion(
	llm,
	model=LLM_MODEL,
	messages=[{"role": "user", "content": prompt}],
	max_tokens=1024,
	temperature=0.25,
	top_p=0.9,
	frequency_penalty=0.1,
	presence_penalty=0.0,
	)

	return response.choices[0].message["content"].strip()


	# ============================================================
	# GRADIO CALLBACKS
	# ============================================================

	def load_pdf(files, state):
	"""
	files: current list from the UI (after add/remove).
	We rebuild state from scratch to stay in sync.

	state = {
	"files": {
	"filename.pdf": {
	"chunks": [...],
	"embeddings": np.array([...])
	},
	...
	},
	"pdf_count": int,
	"total_chunks": int,
	"history": [...],
	"facts": [...]
	}
	"""
	# If user cleared all files
	if not files:
	state["files"] = {}
	state["pdf_count"] = 0
	state["total_chunks"] = 0

	return (
	"No file uploaded.",
	state,
	gr.update(choices=[], value=None),
	)

	try:
	# Reset and rebuild from current files only
	state["files"] = {}
	state["pdf_count"] = 0
	state["total_chunks"] = 0

	# Deduplicate while preserving order
	unique_files = list(dict.fromkeys(files))

	for file in unique_files:
	filename = os.path.basename(str(file))

	text = read_pdf(file)
	chunks = chunk_text(text)

	if not chunks:
	continue

	embeddings = embedding_model.encode(chunks, convert_to_numpy=True)

	state["files"][filename] = {
	"chunks": chunks,
	"embeddings": embeddings,
	}

	state["pdf_count"] += 1
	state["total_chunks"] += len(chunks)

	status_msg = (
	f"{state['pdf_count']} PDF(s) loaded. "
	f"Total chunks: {state['total_chunks']}"
	)

	return (
	status_msg,
	state,
	gr.update(choices=list(state["files"].keys()), value=None),
	)

	except Exception as e:
	return (
	f"Error processing PDF: {e}",
	state,
	gr.update(choices=list(state["files"].keys()), value=None),
	)


	def remove_pdf(filename, state):
	if filename not in state["files"]:
	return (
	f"{filename} not found.",
	state,
	gr.update(choices=list(state["files"].keys()), value=None),
	)

	# Remove the file entry
	state["files"].pop(filename, None)

	# Recompute counts
	pdf_count = len(state["files"])
	total_chunks = sum(len(data["chunks"]) for data in state["files"].values())

	state["pdf_count"] = pdf_count
	state["total_chunks"] = total_chunks

	status_msg = (
	f"Removed {filename}. "
	f"Total PDFs: {pdf_count}. Total chunks: {total_chunks}"
	)

	return (
	status_msg,
	state,
	gr.update(choices=list(state["files"].keys()), value=None),
	)


	def answer_question(question, state):
	print("HF_TOKEN loaded:", HF_TOKEN is not None)

	if HF_TOKEN:
	print("HF_TOKEN length:", len(HF_TOKEN))
	print("HF_TOKEN preview:", HF_TOKEN[:6] + "..." + HF_TOKEN[-4:])
	else:
	print("HF_TOKEN is None or empty")

	question = question.strip()
	if not question:
	return "Please enter a question."

	if not state["files"]:
	return "Please upload at least one PDF first."

	# Simple persistent memory hook: "Remember that ..."
	lower_q = question.lower()
	if lower_q.startswith("remember that"):
	fact = question[len("remember that"):].strip(" .")
	if fact:
	state["facts"].append(fact)
	return f"Okay, I will remember that {fact}."
	else:
	return "What should I remember?"

	# RAG context (per-file balanced retrieval)
	rag_context = retrieve_context_per_file(question, state, k=10)

	# Conversation memory (last 5 turns)
	memory_context = ""
	if state["history"]:
	memory_context = "\n\n".join(
	[
	f"User: {h['user']}\nAssistant: {h['assistant']}"
	for h in state["history"][-5:]
	]
	)

	# Persistent facts
	facts_context = ""
	if state["facts"]:
	facts_context = "User Facts:\n" + "\n".join(state["facts"]) + "\n\n"

	# Merge all contexts
	full_context_parts = []
	if facts_context:
	full_context_parts.append(facts_context)
	if memory_context:
	full_context_parts.append("Conversation Memory:\n" + memory_context)
	if rag_context:
	full_context_parts.append("Document Context:\n" + rag_context)

	full_context = "\n\n".join(full_context_parts)

	# Generate answer
	answer = generate_answer(question, full_context)

	# Save conversational memory
	state["history"].append({"user": question, "assistant": answer})

	return answer


	# ============================================================
	# GRADIO UI
	# ============================================================

	with gr.Blocks(title="PDF Q&A Chatbot") as demo:
	gr.Markdown("## 📄 PDF Q&A Chatbot\nUpload PDFs and ask questions about their content.")

	state = gr.State(
	{
	"files": {}, # filename -> {chunks, embeddings}
	"pdf_count": 0,
	"total_chunks": 0,
	"history": [], # conversational memory
	"facts": [], # persistent user facts
	}
	)

	with gr.Row():
	# LEFT COLUMN — PDF upload + removal
	with gr.Column(scale=1):
	pdf_input = gr.File(
	label="Upload PDF(s)",
	file_types=[".pdf"],
	file_count="multiple", # proper multi-upload in Gradio 6
	)
	pdf_status = gr.Textbox(label="Status", interactive=False)

	file_selector = gr.Dropdown(
	label="Remove a PDF",
	choices=[],
	interactive=True,
	)
	remove_button = gr.Button("Remove Selected PDF")

	# RIGHT COLUMN — question + clear + answer
	with gr.Column(scale=2):
	with gr.Row():
	question = gr.Textbox(
	label="Your question",
	lines=2,
	scale=3,
	)
	clear_button = gr.Button(
	"Clear",
	scale=1,
	min_width=80,
	elem_classes="clear-btn",
	)

	answer = gr.Textbox(label="Answer", lines=8)
	ask_button = gr.Button("Ask")

	# Bind events AFTER components exist
	pdf_input.change(
	load_pdf,
	inputs=[pdf_input, state],
	outputs=[pdf_status, state, file_selector],
	)

	remove_button.click(
	remove_pdf,
	inputs=[file_selector, state],
	outputs=[pdf_status, state, file_selector],
	)

	ask_button.click(
	answer_question,
	inputs=[question, state],
	outputs=answer,
	)



	clear_button.click(lambda: "", outputs=question)


	if __name__ == "__main__":
	demo.launch(
	css=".clear-btn button { height: 100%; min-height: 52px; }"
	)