Upload folder using huggingface_hub

README.md

@@ -37,6 +37,11 @@ or prediction engine — this is Garrett's own tooling, unrelated to GSE.**
 - Streaming responses
 - Copy button on every reply (hover a message)
 - Chat history capped at the last 12 turns to bound context growth
+- **Documents tab**: upload PDF/DOCX/PPTX/HTML/TXT/MD, parsed locally via
+  `docling` (MIT, no external API), chunked, and indexed with BM25 keyword
+  search (`rank_bm25`, pure Python, no embedding model, no extra GPU
+  weight). Check "Search my documents" in the Chat tab to have the model
+  retrieve and cite relevant excerpts instead of answering blind.
 
 ## Swapping the model
 

app.py

@@ -9,9 +9,13 @@ from transformers import (
     BitsAndBytesConfig,
     TextIteratorStreamer,
 )
+from docling.document_converter import DocumentConverter
+from rank_bm25 import BM25Okapi
 
 MODEL_ID = "Qwen/Qwen2.5-72B-Instruct"
 MAX_HISTORY_TURNS = 12
+CHUNK_CHARS = 800
+TOP_K = 4
 
 SYSTEM_PROMPTS = {
     "General": "You are a helpful assistant.",
@@ -43,16 +47,82 @@ model = AutoModelForCausalLM.from_pretrained(
     attn_implementation="sdpa",
 )
 
-
-def _duration(message, history, system_prompt_name, max_new_tokens, temperature):
+_docling_converter = None
+CHUNKS = []  # list of {"source": filename, "text": chunk_text}
+BM25_INDEX = None
+
+
+def _docling():
+    global _docling_converter
+    if _docling_converter is None:
+        _docling_converter = DocumentConverter()
+    return _docling_converter
+
+
+def _chunk_markdown(md_text, source, max_chars=CHUNK_CHARS):
+    paragraphs = [p.strip() for p in md_text.split("\n\n") if p.strip()]
+    chunks, current = [], ""
+    for p in paragraphs:
+        if len(current) + len(p) + 1 <= max_chars:
+            current = f"{current}\n{p}".strip()
+        else:
+            if current:
+                chunks.append(current)
+            current = p
+    if current:
+        chunks.append(current)
+    return [{"source": source, "text": c} for c in chunks]
+
+
+def index_documents(files):
+    global CHUNKS, BM25_INDEX
+    if not files:
+        return "No files given.", ""
+    new_chunks = []
+    for f in files:
+        result = _docling().convert(f.name if hasattr(f, "name") else f)
+        md = result.document.export_to_markdown()
+        source = (f.name if hasattr(f, "name") else str(f)).split("/")[-1]
+        new_chunks.extend(_chunk_markdown(md, source))
+    CHUNKS = new_chunks  # replace — one working set at a time, kept simple
+    tokenized = [c["text"].lower().split() for c in CHUNKS]
+    BM25_INDEX = BM25Okapi(tokenized) if tokenized else None
+    sources = sorted({c["source"] for c in CHUNKS})
+    status = f"Indexed {len(sources)} document(s), {len(CHUNKS)} chunks."
+    listing = "\n".join(f"- {s}" for s in sources)
+    return status, listing
+
+
+def _retrieve(query, k=TOP_K):
+    if BM25_INDEX is None or not CHUNKS:
+        return []
+    scores = BM25_INDEX.get_scores(query.lower().split())
+    ranked = sorted(range(len(CHUNKS)), key=lambda i: scores[i], reverse=True)
+    return [CHUNKS[i] for i in ranked[:k] if scores[i] > 0]
+
+
+def _duration(message, history, system_prompt_name, use_docs, max_new_tokens, temperature):
     return min(120, 40 + int(max_new_tokens) // 10)
 
 
 @spaces.GPU(duration=_duration)
-def generate(message, history, system_prompt_name, max_new_tokens, temperature):
+def generate(message, history, system_prompt_name, use_docs, max_new_tokens, temperature):
     history = history[-(MAX_HISTORY_TURNS * 2):]
     system_prompt = SYSTEM_PROMPTS.get(system_prompt_name, SYSTEM_PROMPTS["General"])
 
+    retrieved = _retrieve(message) if use_docs else []
+    if retrieved:
+        context_block = "\n\n".join(
+            f"[Source: {c['source']}]\n{c['text']}" for c in retrieved
+        )
+        system_prompt = (
+            f"{system_prompt}\n\nUse the following excerpts from the user's own "
+            "uploaded documents if relevant to their question. Cite the source "
+            "filename in brackets when you use one. If the excerpts don't answer "
+            "the question, say so and answer from general knowledge instead.\n\n"
+            f"{context_block}"
+        )
+
     messages = [{"role": "system", "content": system_prompt}]
     for turn in history:
         messages.append({"role": turn["role"], "content": turn["content"]})
@@ -90,23 +160,42 @@ def generate(message, history, system_prompt_name, max_new_tokens, temperature):
 
 chatbot = gr.Chatbot(type="messages", show_copy_button=True, label="Chat")
 
-demo = gr.ChatInterface(
-    fn=generate,
-    type="messages",
-    chatbot=chatbot,
-    additional_inputs=[
-        gr.Dropdown(
-            list(SYSTEM_PROMPTS.keys()), value="General", label="System prompt"
-        ),
-        gr.Slider(64, 1024, value=512, step=64, label="Max new tokens"),
-        gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature"),
-    ],
-    title="Beexly Studio Chat (ZeroGPU)",
-    description=(
+with gr.Blocks(title="Beexly Studio Chat (ZeroGPU)") as demo:
+    gr.Markdown(
+        "# Beexly Studio Chat (ZeroGPU)\n"
         f"Personal free-GPU chat — `{MODEL_ID}` (4-bit), streaming, on ZeroGPU. "
-        "Hover any reply for a copy button. Not part of any Galaxy Sports Edge product."
-    ),
-)
+        "Not part of any Galaxy Sports Edge product."
+    )
+    with gr.Tab("Chat"):
+        chat_iface = gr.ChatInterface(
+            fn=generate,
+            type="messages",
+            chatbot=chatbot,
+            additional_inputs=[
+                gr.Dropdown(
+                    list(SYSTEM_PROMPTS.keys()), value="General", label="System prompt"
+                ),
+                gr.Checkbox(
+                    label="Search my documents (Documents tab)", value=False
+                ),
+                gr.Slider(64, 1024, value=512, step=64, label="Max new tokens"),
+                gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature"),
+            ],
+        )
+    with gr.Tab("Documents"):
+        gr.Markdown(
+            "Upload documents (PDF/DOCX/PPTX/HTML/TXT/MD), parsed locally via "
+            "`docling` (MIT, no external API) and retrieved with BM25 keyword "
+            "search — no embedding model, no extra GPU weight. Check "
+            "'Search my documents' in the Chat tab to use them."
+        )
+        docs_files = gr.File(label="Documents", file_count="multiple")
+        docs_btn = gr.Button("Index documents")
+        docs_status = gr.Textbox(label="Status", interactive=False)
+        docs_list = gr.Textbox(label="Indexed sources", interactive=False)
+        docs_btn.click(
+            index_documents, inputs=docs_files, outputs=[docs_status, docs_list]
+        )
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -4,3 +4,5 @@ transformers>=4.46.0
 accelerate>=1.1.0
 torch>=2.4.0
 bitsandbytes>=0.44.0
+docling
+rank_bm25