Upload 3 files

README.md

@@ -0,0 +1,13 @@
+---
+title: MarisTest
+emoji: 🏢
+colorFrom: red
+colorTo: gray
+sdk: gradio
+sdk_version: 5.44.1
+app_file: app.py
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,464 @@
+import os
+import io
+import json
+import pathlib
+import shutil
+from typing import List, Tuple, Dict
+
+import gradio as gr
+import numpy as np
+import faiss
+from sentence_transformers import SentenceTransformer
+from pypdf import PdfReader
+import fitz # PyMuPDF
+from collections import defaultdict
+from openai import OpenAI
+
+# =========================
+# LLM Endpoint (NVIDIA Integrate via OpenAI SDK; non-streaming)
+# =========================
+
+API_KEY = os.environ.get("NVAPI_KEY")
+
+if not API_KEY:
+    raise RuntimeError("Missing NVAPI_KEY (set it in Hugging Face: Settings → Variables and secrets).")
+
+client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=API_KEY)
+
+MODEL_NAME = "openai/gpt-oss-120b:free"
+GEN_TEMPERATURE = 0.2
+GEN_TOP_P = 0.95
+GEN_MAX_TOKENS = 1024
+
+# =========================
+# Vector Retrieval/Persistence (prefer /data; otherwise ./store)
+# =========================
+
+EMB_MODEL_NAME = "intfloat/multilingual-e5-base"  # Supports both Chinese and English
+
+def choose_store_dir() -> Tuple[str, bool]:
+    """Return (store_dir, is_persistent). Prefer /data/rag_store if writable."""
+    data_root = "/data"
+    if os.path.isdir(data_root) and os.access(data_root, os.W_OK):
+        d = os.path.join(data_root, "rag_store")
+        try:
+            os.makedirs(d, exist_ok=True)
+            testf = os.path.join(d, ".write_test")
+            with open(testf, "w", encoding="utf-8") as f:
+                f.write("ok")
+            os.remove(testf)
+            return d, True
+        except Exception:
+            pass
+    
+    d = os.path.join(os.getcwd(), "store")
+    os.makedirs(d, exist_ok=True)
+    return d, False
+
+STORE_DIR, IS_PERSISTENT = choose_store_dir()
+META_PATH = os.path.join(STORE_DIR, "meta.json")
+INDEX_PATH = os.path.join(STORE_DIR, "faiss.index")
+
+# One-time migration from old ./store to /data/rag_store (if not yet exists)
+LEGACY_STORE_DIR = os.path.join(os.getcwd(), "store")
+
+def migrate_legacy_if_any():
+    try:
+        if IS_PERSISTENT:
+            legacy_meta = os.path.join(LEGACY_STORE_DIR, "meta.json")
+            legacy_index = os.path.join(LEGACY_STORE_DIR, "faiss.index")
+            
+            if (not os.path.exists(META_PATH) or not os.path.exists(INDEX_PATH)) \
+                and os.path.isdir(LEGACY_STORE_DIR) \
+                and os.path.exists(legacy_meta) and os.path.exists(legacy_index):
+                shutil.copyfile(legacy_meta, META_PATH)
+                shutil.copyfile(legacy_index, INDEX_PATH)
+    except Exception:
+        pass
+
+migrate_legacy_if_any()
+
+_emb_model = None
+_index: faiss.Index = None
+_meta: Dict[str, Dict] = {}
+
+# ========== Adjustable Parameter Defaults ==========
+DEFAULT_TOP_K = 6
+DEFAULT_POOL_K = 40
+DEFAULT_PER_SOURCE_CAP = 2
+DEFAULT_STRATEGY = "mmr"  # "mmr" or "round_robin"
+DEFAULT_MMR_LAMBDA = 0.5
+
+# ---------- Basic Tools ----------
+
+def get_emb_model():
+    global _emb_model
+    if _emb_model is None:
+        _emb_model = SentenceTransformer(EMB_MODEL_NAME)
+    return _emb_model
+
+def _ensure_index(dim: int):
+    global _index
+    if _index is None:
+        _index = faiss.IndexFlatIP(dim)  # Normalize vectors first → inner product = cosine
+
+def _persist():
+    faiss.write_index(_index, INDEX_PATH)
+    with open(META_PATH, "w", encoding="utf-8") as f:
+        json.dump(_meta, f, ensure_ascii=False)
+
+def _load_if_any():
+    global _index, _meta
+    if os.path.exists(INDEX_PATH) and os.path.exists(META_PATH):
+        _index = faiss.read_index(INDEX_PATH)
+        with open(META_PATH, "r", encoding="utf-8") as f:
+            _meta = json.load(f)
+
+def _chunk_text(text: str, chunk_size: int = 800, overlap: int = 120) -> List[str]:
+    text = text.replace("\u0000", "")
+    res, i, n = [], 0, len(text)
+    while i < n:
+        j = min(i + chunk_size, n)
+        seg = text[i:j].strip()
+        if seg:
+            res.append(seg)
+        i = max(0, j - overlap)
+        if j >= n:
+            break
+    return res
+
+# ---------- Robust File Reading ----------
+
+def _read_bytes(file) -> bytes:
+    if isinstance(file, dict):
+        p = file.get("path") or file.get("name")
+        if p and os.path.exists(p):
+            with open(p, "rb") as f:
+                return f.read()
+        if "data" in file and isinstance(file["data"], (bytes, bytearray)):
+            return bytes(file["data"])
+    
+    if isinstance(file, (str, pathlib.Path)):
+        with open(file, "rb") as f:
+            return f.read()
+    
+    if hasattr(file, "read"):
+        try:
+            if hasattr(file, "seek"):
+                try:
+                    file.seek(0)
+                except Exception:
+                    pass
+            return file.read()
+        finally:
+            try:
+                file.close()
+            except Exception:
+                pass
+    
+    raise ValueError("Unsupported file type from gr.File")
+
+def _decode_best_effort(raw: bytes) -> str:
+    for enc in ["utf-8", "cp932", "shift_jis", "cp950", "big5", "gb18030", "latin-1"]:
+        try:
+            return raw.decode(enc)
+        except Exception:
+            continue
+    return raw.decode("utf-8", errors="ignore")
+
+def _read_pdf(file_bytes: bytes) -> str:
+    # 1) PyMuPDF
+    try:
+        with fitz.open(stream=file_bytes, filetype="pdf") as doc:
+            if doc.is_encrypted:
+                try:
+                    doc.authenticate("")  # Try empty password
+                except Exception:
+                    pass
+            texts = [(page.get_text("text") or "") for page in doc]
+            txt = "\n".join(texts)
+            if txt.strip():
+                return txt
+    except Exception:
+        pass
+    
+    # 2) fallback: pypdf
+    try:
+        reader = PdfReader(io.BytesIO(file_bytes))
+        pages = []
+        for p in reader.pages:
+            try:
+                pages.append(p.extract_text() or "")
+            except Exception:
+                pages.append("")
+        return "\n".join(pages)
+    except Exception:
+        return ""
+
+def _read_any(file) -> str:
+    if isinstance(file, dict):
+        name = (file.get("orig_name") or file.get("name") or file.get("path") or "upload").lower()
+    else:
+        name = getattr(file, "name", None) or (str(file) if isinstance(file, (str, pathlib.Path)) else "upload")
+    name = name.lower()
+    
+    raw = _read_bytes(file)
+    if name.endswith(".pdf"):
+        return _read_pdf(raw).replace("\u0000", "")
+    return _decode_best_effort(raw).replace("\u0000", "")
+
+# ---------- Build Corpus ----------
+
+def build_corpus(files) -> str:
+    if not files:
+        return "No files selected."
+    
+    emb_model = get_emb_model()
+    chunks, sources, failed = [], [], []
+    total_chars = 0
+    
+    for f in files:
+        if isinstance(f, dict):
+            fname = f.get("orig_name") or f.get("name") or f.get("path") or "uploaded"
+        else:
+            fname = getattr(f, "name", None) or (os.path.basename(f) if isinstance(f, (str, pathlib.Path)) else "uploaded")
+        
+        try:
+            text = _read_any(f) or ""
+            total_chars += len(text)
+            parts = _chunk_text(text)
+            if not parts:
+                failed.append(fname)
+                continue
+            chunks.extend(parts)
+            sources.extend([fname] * len(parts))
+        except Exception as e:
+            failed.append(f"{fname} (err: {e})")
+    
+    if not chunks:
+        tier = "Persistent (/data)" if IS_PERSISTENT else "Ephemeral (./store)"
+        return f"No text extracted (please check file type/encoding; read {total_chars} characters this time).\nCurrent storage path: {STORE_DIR} [{tier}]"
+    
+    passages = [f"passage: {c}" for c in chunks]  # e5 prefix
+    vec = emb_model.encode(passages, batch_size=64, convert_to_numpy=True, normalize_embeddings=True)
+    
+    _ensure_index(vec.shape[1])
+    _index.add(vec)
+    
+    base = len(_meta)
+    for i, (src, c) in enumerate(zip(sources, chunks)):
+        _meta[str(base + i)] = {"source": src, "text": c}
+    
+    _persist()
+    
+    msg = f"Indexing complete: added {len(chunks)} chunks; current total chunks in corpus ≈ {_index.ntotal}."
+    if failed:
+        preview = ", ".join(failed[:5])
+        more = "" if len(failed) <= 5 else f" (and {len(failed)-5} more not shown)"
+        msg += f"\nNote: {len(failed)} files failed to extract text or were empty: {preview}{more}"
+    
+    tier = "Persistent (/data)" if IS_PERSISTENT else "Ephemeral (./store)"
+    msg += f"\nCurrent storage path: {STORE_DIR} [{tier}]"
+    
+    return msg
+
+# ---------- Retrieval (Candidate Pool) ----------
+
+def _encode_query_vec(query: str) -> np.ndarray:
+    return get_emb_model().encode([f"query: {query}"], convert_to_numpy=True, normalize_embeddings=True)
+
+def retrieve_candidates(qvec: np.ndarray, pool_k: int = DEFAULT_POOL_K) -> List[Tuple[str, float]]:
+    if _index is None or _index.ntotal == 0:
+        return []
+    
+    pool_k = min(pool_k, _index.ntotal)
+    D, I = _index.search(qvec, pool_k)
+    return [(str(idx), float(score)) for idx, score in zip(I[0], D[0]) if idx != -1]
+
+# ---------- Diversification Strategies ----------
+
+def select_diverse_by_source(cands: List[Tuple[str, float]], top_k: int = DEFAULT_TOP_K, per_source_cap: int = DEFAULT_PER_SOURCE_CAP) -> List[Tuple[str, float]]:
+    if not cands:
+        return []
+    
+    by_src: Dict[str, List[Tuple[str, float]]] = defaultdict(list)
+    for cid, s in cands:
+        m = _meta.get(cid)
+        if not m:
+            continue
+        by_src[m["source"]].append((cid, s))
+    
+    for src in by_src:
+        by_src[src] = by_src[src][:per_source_cap]
+    
+    picked, src_items, ptrs = [], [(s, it) for s, it in by_src.items()], {s: 0 for s in by_src}
+    while len(picked) < top_k:
+        advanced = False
+        for src, items in src_items:
+            i = ptrs[src]
+            if i < len(items):
+                picked.append(items[i])
+                ptrs[src] = i + 1
+                advanced = True
+                if len(picked) >= top_k:
+                    break
+        if not advanced:
+            break
+    
+    if len(picked) < top_k:
+        seen = {cid for cid, _ in picked}
+        for cid, s in cands:
+            if cid not in seen:
+                picked.append((cid, s))
+                seen.add(cid)
+                if len(picked) >= top_k:
+                    break
+    
+    return picked[:top_k]
+
+def _encode_chunks_text(cids: List[str]) -> np.ndarray:
+    texts = [f"passage: {(_meta.get(cid) or {}).get('text','')}" for cid in cids]
+    return get_emb_model().encode(texts, convert_to_numpy=True, normalize_embeddings=True)
+
+def select_diverse_mmr(cands: List[Tuple[str, float]], qvec: np.ndarray, top_k: int = DEFAULT_TOP_K, mmr_lambda: float = DEFAULT_MMR_LAMBDA) -> List[Tuple[str, float]]:
+    if not cands:
+        return []
+    
+    cids = [cid for cid, _ in cands]
+    cvecs = _encode_chunks_text(cids)
+    sim_to_q = (cvecs @ qvec.T).reshape(-1)
+    
+    selected, remaining = [], set(range(len(cids)))
+    while len(selected) < min(top_k, len(cids)):
+        if not selected:
+            i = int(np.argmax(sim_to_q))
+            selected.append(i)
+            remaining.remove(i)
+            continue
+        
+        S = cvecs[selected]
+        sim_to_S = (cvecs[list(remaining)] @ S.T)
+        max_sim_to_S = sim_to_S.max(axis=1) if sim_to_S.size > 0 else np.zeros((len(remaining),), dtype=np.float32)
+        
+        sim_q_rem = sim_to_q[list(remaining)]
+        mmr_scores = mmr_lambda * sim_q_rem - (1.0 - mmr_lambda) * max_sim_to_S
+        
+        j_rel = int(np.argmax(mmr_scores))
+        j = list(remaining)[j_rel]
+        selected.append(j)
+        remaining.remove(j)
+    
+    return [(cids[i], float(sim_to_q[i])) for i in selected][:top_k]
+
+def retrieve_diverse(query: str,
+                    top_k: int = DEFAULT_TOP_K,
+                    pool_k: int = DEFAULT_POOL_K,
+                    per_source_cap: int = DEFAULT_PER_SOURCE_CAP,
+                    strategy: str = DEFAULT_STRATEGY,
+                    mmr_lambda: float = DEFAULT_MMR_LAMBDA) -> List[Tuple[str, float]]:
+    qvec = _encode_query_vec(query)
+    cands = retrieve_candidates(qvec, pool_k=pool_k)
+    
+    if strategy == "mmr":
+        return select_diverse_mmr(cands, qvec, top_k=top_k, mmr_lambda=mmr_lambda)
+    return select_diverse_by_source(cands, top_k=top_k, per_source_cap=per_source_cap)
+
+# ---------- Format Chunks ----------
+
+def _format_ctx(hits: List[Tuple[str, float]]) -> str:
+    if not hits:
+        return ""
+    
+    lines = []
+    for cid, _ in hits:
+        m = _meta.get(cid)
+        if not m:
+            continue
+        source_clean = m.get("source", "")
+        text_clean = (m.get("text", "") or "").replace("\n", " ")
+        lines.append(f"[{cid}] ({source_clean}) " + text_clean)
+    
+    return "\n".join(lines[:10])
+
+# =========================
+# LLM Conversation (non-streaming → avoid StopAsyncIteration)
+# =========================
+
+def chat_fn(message, history, strategy, top_k, pool_k, per_source_cap, mmr_lambda):
+    hits = retrieve_diverse(
+        message,
+        top_k=int(top_k),
+        pool_k=int(pool_k),
+        per_source_cap=int(per_source_cap),
+        strategy=str(strategy),
+        mmr_lambda=float(mmr_lambda),
+    )
+    
+    ctx = _format_ctx(hits) if hits else "(Current index is empty or no matching chunks found)"
+    
+    sys_blocks = [
+        "You are a rigorous real estate market research assistant. Your answers must be based on retrieved content with evidence and source numbers cited. If retrieval is insufficient, please clearly explain the shortcomings.",
+        f"Below are the available reference chunks (with numbers and sources). When answering, please cite the numbers, e.g., [3].\n\n{ctx}",
+    ]
+    
+    messages = [{"role": "system", "content": "\n\n".join(sys_blocks)}]
+    for u, a in history:
+        messages.append({"role": "user", "content": u})
+        messages.append({"role": "assistant", "content": a})
+    messages.append({"role": "user", "content": message})
+    
+    try:
+        resp = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            temperature=GEN_TEMPERATURE,
+            top_p=GEN_TOP_P,
+            max_tokens=GEN_MAX_TOKENS,
+            stream=False,  # Non-streaming, most stable
+        )
+        return resp.choices[0].message.content
+    except Exception as e:
+        return f"[Exception] {repr(e)}"
+
+# =========================
+# Gradio Interface (Slider/Dropdown takes effect immediately)
+# =========================
+
+with gr.Blocks(title="RAG (Candidate Pool + Diversification) | NVIDIA Integrate (Non-streaming)") as demo:
+    tier = "Persistent `/data`" if IS_PERSISTENT else "Ephemeral `./store` (may be lost on restart)"
+    gr.Markdown(
+        f"### RAG (Plain text, no OCR)\n"
+        f"- Current storage path: `{STORE_DIR}` [{tier}]\n"
+        "- Full corpus retrieval + MMR / Round-Robin diversification; chunks are labeled with source filename and chunk id.\n"
+        f"- LLM: `{MODEL_NAME}` (non-streaming, avoids StopAsyncIteration)."
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            files = gr.File(file_count="multiple", file_types=[".pdf", ".txt"], label="Upload Corpus (multiple files allowed)")
+            out = gr.Textbox(label="Index Status", interactive=False)
+            gr.Button("Build/Update Index", variant="primary").click(build_corpus, inputs=files, outputs=out)
+            
+            gr.Markdown("#### Retrieval Parameters (adjustable based on data volume)")
+            strategy = gr.Dropdown(choices=["mmr", "round_robin"], value=DEFAULT_STRATEGY, label="Diversification Strategy")
+            top_k = gr.Slider(1, 12, value=DEFAULT_TOP_K, step=1, label="Number of chunks for model (top_k)")
+            pool_k = gr.Slider(10, 200, value=DEFAULT_POOL_K, step=5, label="Candidate pool size (pool_k)")
+            per_source_cap = gr.Slider(1, 5, value=DEFAULT_PER_SOURCE_CAP, step=1, label="Per-source limit (for round_robin)")
+            mmr_lambda = gr.Slider(0.0, 1.0, value=DEFAULT_MMR_LAMBDA, step=0.05, label="MMR λ (higher = closer to query)")
+        
+        with gr.Column(scale=2):
+            gr.ChatInterface(
+                fn=chat_fn,  # Non-streaming function
+                additional_inputs=[strategy, top_k, pool_k, per_source_cap, mmr_lambda],
+                title="Chat (RAG)",
+                description=f"First retrieve from full corpus → diversify chunk selection → generate answer with {MODEL_NAME}.",
+            )
+
+# Cold start: load existing index
+try:
+    _load_if_any()
+except Exception:
+    pass
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=4.0.0
+openai>=1.40.0
+faiss-cpu>=1.7.4
+sentence-transformers>=2.7.0
+pypdf>=4.2.0
+pymupdf>=1.24.9
+numpy>=1.26.0