File size: 11,695 Bytes
451bc89
2cc309f
f3aa49e
1c18139
b1e4a16
 
4a597f7
451bc89
 
2cc309f
4a597f7
 
 
2cc309f
9766920
2cc309f
451bc89
 
 
 
 
 
 
 
 
 
 
 
 
 
1a03e36
4a597f7
 
c30140e
451bc89
 
 
b1e4a16
f3aa49e
451bc89
f3aa49e
 
451bc89
 
 
 
 
 
 
 
 
 
 
 
f3aa49e
 
4a597f7
 
 
b1e4a16
4a597f7
b1e4a16
4a597f7
b1e4a16
 
 
4a597f7
 
b1e4a16
 
4a597f7
b1e4a16
 
 
4a597f7
b1e4a16
 
 
 
 
4a597f7
b1e4a16
 
 
0135e8c
 
 
b1e4a16
0135e8c
451bc89
 
 
 
0135e8c
 
4a597f7
0135e8c
 
4a597f7
0135e8c
 
 
b1e4a16
0135e8c
 
b1e4a16
451bc89
0135e8c
451bc89
 
0135e8c
b1e4a16
451bc89
0135e8c
 
b1e4a16
0135e8c
 
 
 
4a597f7
b1e4a16
4a597f7
f3aa49e
4a597f7
b1e4a16
4a597f7
 
 
 
 
 
 
b1e4a16
 
 
4a597f7
b1e4a16
 
 
451bc89
 
 
 
 
 
 
 
 
 
 
 
2cc309f
 
4a597f7
 
 
2cc309f
bbe2a51
451bc89
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bbe2a51
5f6a59c
 
0135e8c
5f6a59c
 
 
 
0135e8c
5f6a59c
b1e4a16
 
451bc89
5f6a59c
 
0135e8c
5f6a59c
451bc89
180b29c
8051cd5
5299745
5f6a59c
5299745
180b29c
0135e8c
2cc309f
0135e8c
bbe2a51
2cc309f
0135e8c
 
 
 
 
 
 
bbe2a51
0135e8c
2cc309f
0135e8c
 
 
 
2cc309f
bbe2a51
0135e8c
bbe2a51
0135e8c
bbe2a51
2cc309f
b1e4a16
8051cd5
 
 
0135e8c
8051cd5
0135e8c
 
765eae2
 
5f6a59c
 
 
0135e8c
5f6a59c
 
451bc89
5f6a59c
2cc309f
451bc89
0135e8c
 
765eae2
0135e8c
 
765eae2
0135e8c
 
 
 
765eae2
 
0135e8c
765eae2
0135e8c
765eae2
bbe2a51
0135e8c
b1e4a16
451bc89
 
 
 
 
 
 
 
dc063ef
 
b1e4a16
 
0135e8c
 
2cc309f
451bc89
dc063ef
 
 
 
 
 
 
 
 
451bc89
dc063ef
 
451bc89
dc063ef
 
 
 
 
 
 
 
 
451bc89
dc063ef
 
 
 
451bc89
dc063ef
 
 
 
 
 
 
 
 
 
451bc89
dc063ef
 
451bc89
dc063ef
 
1c18139
 
 
4a597f7
 
 
1c18139
2c7e565
 
b1e4a16
0135e8c
 
451bc89
0135e8c
 
451bc89
 
0135e8c
 
1c18139
6cfe0b8
451bc89
6cfe0b8
0135e8c
 
 
451bc89
0135e8c
6cfe0b8
8dfb278
340d4e8
 
 
0135e8c
340d4e8
2c7e565
 
451bc89
6cfe0b8
8594f0e
340d4e8
 
 
0135e8c
340d4e8
 
 
 
 
0135e8c
340d4e8
2c7e565
6cfe0b8
 
8dfb278
451bc89
0135e8c
 
 
 
 
40024c6
340d4e8
 
 
0135e8c
340d4e8
2cc309f
340d4e8
 
 
0135e8c
340d4e8
7f73068
451bc89
 
340d4e8
10205ee
bbe2a51
b1e4a16
340d4e8
0135e8c
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410

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; }"
    )