Initial commit (Xet)

.gitattributes

@@ -0,0 +1 @@
+*.npy filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,96 @@
+file
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheelhouse/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Type checkers
+.mypy_cache/
+.pyre/
+.pytype/
+
+# Linting / formatting
+.ruff_cache/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# Django / Flask / FastAPI (common)
+*.log
+local_settings.py
+instance/
+.webassets-cache
+
+# Sphinx documentation
+docs/_build/
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+env.bak/
+venv.bak/
+
+# Environment variables / secrets
+.env
+.env.*
+*.env
+
+# Editors / IDEs
+.vscode/
+.idea/
+*.sublime-project
+*.sublime-workspace
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# PyCharm
+*.iml

README.md

@@ -0,0 +1,83 @@
+---
+title: RAG QA (BM25 + Dense + Reranker)
+emoji: 🔎
+colorFrom: indigo
+colorTo: pink
+sdk: gradio
+sdk_version: "6.1.0"
+python_version: "3.10"
+app_file: app.py
+pinned: false
+short_description: RAG Question Answering demo with BM25 + Dense retrieval and optional reranker.
+---
+
+# RAG QA (BM25 + Dense + OpenAI-compatible providers)
+
+Це навчальний проєкт **Retrieval-Augmented Generation (RAG)** для задачі **Question Answering**: система знаходить релевантні уривки в корпусі документів і відповідає на питання, спираючись лише на знайдений контекст. :contentReference[oaicite:2]{index=2}
+
+## Demo
+- Hugging Face Space: *(посилання після деплою)*  
+- GitHub: https://github.com/DunasAnastasiia/llm-chat-project
+
+---
+
+## Як працює система (архітектура)
+
+**Джерело даних (Dataset)**  
+- Hugging Face dataset: `rag-datasets/rag-mini-wikipedia` :contentReference[oaicite:3]{index=3}
+
+**Chunking**  
+- Простий character-based chunking з overlap: `chunk_chars=900`, `overlap_chars=150` 
+
+**Retriever**
+- **BM25** (лексичний пошук по ключових словах)
+- **Dense retrieval** (Sentence-Transformers embeddings)
+- Можна **вмикати/вимикати** BM25 та Dense окремо (для порівняння якості). 
+
+**Reranker (optional)**
+- Cross-encoder `cross-encoder/ms-marco-MiniLM-L-6-v2` для реранкінгу кандидатів (опціонально). 
+
+**LLM (Generation)**
+- Використовується **OpenAI-compatible Chat Completions API** через бібліотеку `openai` із параметром `base_url`.
+- Працює з провайдерами типу **Groq** / **OpenRouter** (ключ вводиться в UI). 
+
+**Citations**
+- В промпті LLM проситься цитувати чанки як `[1] [2] ...`, а в UI показується список retrieved chunks з `source_id` / `chunk_id`. 
+
+**UI**
+- Gradio (web app), тумблери `Use BM25 / Use Dense / Use Reranker`, поле для API key. :contentReference[oaicite:9]{index=9}
+
+## Приклади запитів: де BM25 кращий, а де Dense
+
+### Запити, де краще справляється BM25
+
+**Q:** Was Abraham Lincoln the first President of the United States?  
+**Очікувана відповідь:** No  
+**Чому BM25 кращий:** запит містить дуже конкретні ключові слова (Abraham Lincoln, first President, United States). BM25 добре працює, коли відповідь лежить у чанку з тими самими словами/формулюванням.
+
+**Q:** Who was the general in charge at the Battle of Antietam?  
+**Очікувана відповідь:** General McClellan  
+**Чому BM25 кращий:** тут є точні “якорі” (Battle of Antietam, general in charge). BM25 зазвичай підтягує уривок, де ці терміни зустрічаються буквально.
+
+---
+
+### Запити, де кращі результати повертає Dense retriever
+
+**Q:** Who assassinated Lincoln?  
+**Очікувана відповідь:** John Wilkes Booth  
+**Чому dense кращий:** типовий приклад “семантичного” запиту — відповідь може бути в уривку, який не повторює точну форму запиту (наприклад, “Lincoln was assassinated by …”). Dense краще ловить перефразування й зв’язки “подія ↔ учасник”.
+
+**Q:** What caused Calvin Jr.'s death?  
+**Очікувана відповідь:** heart attack  
+**Чому dense кращий:** питання може збігатися з текстом не дослівно (наприклад, у корпусі “died of a heart attack”). Dense часто краще дістає такі уривки навіть без точного збігу слів.
+
+
+---
+
+## Швидкий старт локально
+
+### 1) Встановлення
+```bash
+python -m venv .venv
+source .venv/bin/activate  # Windows: .venv\Scripts\activate
+pip install -r requirements.txt

app.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import gradio as gr
+
+from rag.retrieve import Retriever
+from rag.llm import answer_with_provider
+
+
+def ensure_retriever(state):
+    if state is None:
+        state = Retriever()
+    return state
+
+
+def defaults_for_provider(provider_name: str) -> tuple[str, str]:
+    """
+    Returns (base_url, default_model) for a given provider.
+    """
+    if provider_name.startswith("Groq"):
+        return "https://api.groq.com/openai/v1", "llama-3.1-8b-instant"
+    if provider_name.startswith("OpenRouter"):
+        return "https://openrouter.ai/api/v1", "meta-llama/llama-3.1-8b-instruct:free"
+    # fallback
+    return "https://api.groq.com/openai/v1", "llama-3.1-8b-instant"
+
+
+def on_provider_change(provider_name: str):
+    base_url, model = defaults_for_provider(provider_name)
+    return base_url, model
+
+
+def run_qa(
+    provider: str,
+    base_url: str,
+    api_key: str,
+    model: str,
+    question: str,
+    use_bm25: bool,
+    use_dense: bool,
+    use_rerank: bool,
+    state,
+):
+    state = ensure_retriever(state)
+
+    if not question or not question.strip():
+        return "Write a question 🙂", "", state
+
+    # Retrieval toggles
+    chunks = state.retrieve(
+        question,
+        use_bm25=use_bm25,
+        use_dense=use_dense,
+        use_rerank=use_rerank,
+    )
+
+    # Show retrieved context
+    ctx = []
+    for i, c in enumerate(chunks, start=1):
+        ctx.append(
+            f"[{i}] ({c.why}, score={c.score:.4f}) source_id={c.source_id}, chunk_id={c.chunk_id}\n{c.text}"
+        )
+    ctx_text = "\n\n---\n\n".join(ctx) if ctx else "(nothing retrieved)"
+
+    # If both retrievers off => "no retrieval" mode
+    if not use_bm25 and not use_dense:
+        ctx_text = "(retrieval is OFF: the model will answer without any context)"
+        chunks_for_llm = []
+    else:
+        chunks_for_llm = [{"chunk_id": c.chunk_id, "source_id": c.source_id, "text": c.text} for c in chunks]
+
+    if not api_key or not api_key.strip():
+        return f"Paste your {provider} API key first.", ctx_text, state
+
+    # Provider call (OpenAI-compatible Chat Completions)
+    try:
+        ans = answer_with_provider(
+            api_key=api_key.strip(),
+            base_url=(base_url or "").strip(),
+            model=(model or "").strip(),
+            question=question,
+            chunks=chunks_for_llm,
+        )
+    except Exception as e:
+        return f"LLM error: {type(e).__name__}: {e}", ctx_text, state
+
+    return ans, ctx_text, state
+
+
+with gr.Blocks(title="RAG QA (BM25 + Dense + OpenAI-compatible providers)") as demo:
+    gr.Markdown(
+        "# RAG QA (HF dataset + BM25 + Dense)\n"
+        "Use a **free-tier OpenAI-compatible provider** (Groq / OpenRouter).\n"
+        "1) Build index: `python -m rag.index`\n"
+        "2) Run UI: `python app.py`\n"
+    )
+
+    state = gr.State(None)
+
+    provider = gr.Dropdown(
+        ["Groq (free tier)", "OpenRouter (free models)"],
+        value="Groq (free tier)",
+        label="Provider",
+    )
+
+    base_url = gr.Textbox(
+        label="Base URL",
+        value="https://api.groq.com/openai/v1",
+        placeholder="https://api.groq.com/openai/v1",
+    )
+
+    api_key = gr.Textbox(
+        label="API key",
+        type="password",
+        placeholder="paste provider key here",
+    )
+
+    model = gr.Textbox(
+        label="Model",
+        value="llama-3.1-8b-instant",
+    )
+
+    provider.change(
+        fn=on_provider_change,
+        inputs=[provider],
+        outputs=[base_url, model],
+    )
+
+    question = gr.Textbox(label="Question", placeholder="Ask something...", lines=2)
+
+    with gr.Row():
+        use_bm25 = gr.Checkbox(value=True, label="Use BM25")
+        use_dense = gr.Checkbox(value=True, label="Use Dense")
+        use_rerank = gr.Checkbox(value=False, label="Use Reranker (optional)")
+
+    btn = gr.Button("Answer")
+
+    answer = gr.Textbox(label="Answer", lines=8)
+    context = gr.Textbox(label="Retrieved chunks", lines=12)
+
+    btn.click(
+        fn=run_qa,
+        inputs=[provider, base_url, api_key, model, question, use_bm25, use_dense, use_rerank, state],
+        outputs=[answer, context, state],
+    )
+
+
+if __name__ == "__main__":
+    import os
+    port = int(os.getenv("PORT", "7860"))
+    demo.launch(server_name="0.0.0.0", server_port=port, share=False)

artifacts/embeddings.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07ccc5bacea16e81d6a77b1971bbc03110547b7b74e05cf1ff74388dc6b9d940
+size 5689472

rag/chunking.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Chunk:
+    chunk_id: int
+    source_id: str
+    text: str
+
+
+def chunk_text(text: str, chunk_chars: int, overlap_chars: int) -> list[str]:
+    """
+    Simple character-based chunking with overlap.
+    Works for any text without requiring tokenizers.
+    """
+    text = (text or "").strip()
+    if not text:
+        return []
+    if chunk_chars <= 0:
+        return [text]
+
+    out: list[str] = []
+    i = 0
+    n = len(text)
+    step = max(1, chunk_chars - max(0, overlap_chars))
+    while i < n:
+        chunk = text[i : i + chunk_chars].strip()
+        if chunk:
+            out.append(chunk)
+        i += step
+    return out

rag/config.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Settings:
+    # Hugging Face dataset
+    dataset_name: str = "rag-datasets/rag-mini-wikipedia"
+    corpus_config: str = "text-corpus"
+    qa_config: str = "question-answer"
+
+    # Chunking
+    chunk_chars: int = 900
+    overlap_chars: int = 150
+
+    # Retrieval
+    top_k_bm25: int = 8
+    top_k_dense: int = 8
+    top_k_final: int = 6
+
+    # Dense model
+    embed_model: str = "sentence-transformers/all-MiniLM-L6-v2"
+
+    # Optional reranker
+    rerank_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+    rerank_top_n: int = 20  # candidates to rerank
+
+    # OpenAI
+    default_openai_model: str = "gpt-4o-mini"
+
+    # Artifacts
+    artifacts_dir: str = "artifacts"
+    chunks_jsonl: str = "chunks.jsonl"
+    embeddings_npy: str = "embeddings.npy"
+
+
+SETTINGS = Settings()

rag/data.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Iterable, Any
+
+
+def pick_text_column(columns: list[str]) -> str:
+    """
+    Try to robustly choose the text field from various corpus schemas.
+    """
+    candidates = ["text", "content", "document", "passage", "passages", "contents", "wiki_text"]
+    for c in candidates:
+        if c in columns:
+            return c
+    return columns[0]
+
+
+def pick_id_column(columns: list[str]) -> str | None:
+    for c in ["id", "doc_id", "document_id", "passage_id", "pid"]:
+        if c in columns:
+            return c
+    return None
+
+
+def iter_corpus_rows(ds) -> Iterable[dict[str, Any]]:
+    for row in ds:
+        yield row

rag/index.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+from tqdm import tqdm
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+
+from rag.config import SETTINGS
+from rag.chunking import chunk_text
+from rag.data import pick_text_column, pick_id_column
+
+
+def _load_corpus_split():
+    ds = load_dataset(SETTINGS.dataset_name, SETTINGS.corpus_config)
+
+    if hasattr(ds, "keys"):
+        if "train" in ds:
+            return ds["train"]
+        if "passages" in ds:
+            return ds["passages"]
+        if "validation" in ds:
+            return ds["validation"]
+        if "test" in ds:
+            return ds["test"]
+        return ds[list(ds.keys())[0]]
+
+    return ds
+
+
+def build_index(limit_docs: int | None = None) -> None:
+    art = Path(SETTINGS.artifacts_dir)
+    art.mkdir(parents=True, exist_ok=True)
+    chunks_path = art / SETTINGS.chunks_jsonl
+    emb_path = art / SETTINGS.embeddings_npy
+
+    # Load corpus (no trust_remote_code; modern datasets removed/limited it)
+    corpus = _load_corpus_split()
+
+    cols = list(corpus.column_names)
+    text_col = pick_text_column(cols)
+    id_col = pick_id_column(cols)
+
+    chunks: list[dict] = []
+    chunk_id = 0
+
+    n_docs = len(corpus) if limit_docs is None else min(limit_docs, len(corpus))
+
+    for idx in tqdm(range(n_docs), desc="Chunking corpus"):
+        row = corpus[int(idx)]
+        raw = row.get(text_col, "")
+        source_id = str(row.get(id_col, idx)) if id_col else str(idx)
+
+        for part in chunk_text(raw, SETTINGS.chunk_chars, SETTINGS.overlap_chars):
+            chunks.append({"chunk_id": chunk_id, "source_id": source_id, "text": part})
+            chunk_id += 1
+
+    # Save chunks
+    with chunks_path.open("w", encoding="utf-8") as f:
+        for ch in chunks:
+            f.write(json.dumps(ch, ensure_ascii=False) + "\n")
+
+    # Compute embeddings
+    model = SentenceTransformer(SETTINGS.embed_model)
+    texts = [c["text"] for c in chunks]
+    emb = model.encode(
+        texts,
+        batch_size=64,
+        show_progress_bar=True,
+        normalize_embeddings=True,
+    )
+    emb = np.asarray(emb, dtype=np.float32)
+    np.save(emb_path, emb)
+
+    print(f"Saved {len(chunks)} chunks -> {chunks_path}")
+    print(f"Saved embeddings shape={emb.shape} -> {emb_path}")
+
+
+if __name__ == "__main__":
+    build_index()

rag/llm.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from openai import OpenAI
+
+
+def build_prompt(question: str, chunks: list[dict]) -> tuple[str, str]:
+    instructions = (
+        "You are a QA assistant. Answer ONLY using the provided context.\n"
+        "If the answer is not in the context, say you don't know based on the context.\n"
+        "When you use information from a chunk, cite it like [1], [2] matching the chunk numbers.\n"
+        "Be concise."
+    )
+
+    ctx_lines = []
+    for i, ch in enumerate(chunks, start=1):
+        ctx_lines.append(
+            f"[{i}] source_id={ch['source_id']} chunk_id={ch['chunk_id']}\n{ch['text']}\n"
+        )
+
+    input_text = (
+        "CONTEXT:\n"
+        + "\n".join(ctx_lines)
+        + "\nQUESTION:\n"
+        + question.strip()
+        + "\n\nANSWER:"
+    )
+    return instructions, input_text
+
+
+def answer_with_provider(
+    api_key: str,
+    base_url: str,
+    model: str,
+    question: str,
+    chunks: list[dict],
+) -> str:
+    """
+    Works with OpenAI-compatible providers (Groq, OpenRouter, Together, etc.)
+    via Chat Completions API.
+    """
+    client = OpenAI(api_key=api_key, base_url=base_url)
+    instructions, input_text = build_prompt(question, chunks)
+
+    resp = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system", "content": instructions},
+            {"role": "user", "content": input_text},
+        ],
+        temperature=0.2,
+    )
+
+    msg = resp.choices[0].message.content
+    return msg or ""

rag/retrieve.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from rank_bm25 import BM25Okapi
+from sentence_transformers import SentenceTransformer, CrossEncoder
+
+from rag.config import SETTINGS
+
+_WORD = re.compile(r"[A-Za-z0-9']+")
+
+
+def tokenize(text: str) -> list[str]:
+    return _WORD.findall((text or "").lower())
+
+
+@dataclass
+class ChunkRec:
+    chunk_id: int
+    source_id: str
+    text: str
+    score: float
+    why: str  # "bm25", "dense", "rerank"
+
+
+class Retriever:
+    def __init__(self) -> None:
+        art = Path(SETTINGS.artifacts_dir)
+        self.chunks = self._load_chunks(art / SETTINGS.chunks_jsonl)
+        self.emb = np.load(art / SETTINGS.embeddings_npy)
+
+        # BM25
+        tokenized = [tokenize(c["text"]) for c in self.chunks]
+        self.bm25 = BM25Okapi(tokenized)
+
+        # Dense encoder
+        self.embedder = SentenceTransformer(SETTINGS.embed_model)
+
+        # Reranker (lazy)
+        self._reranker: CrossEncoder | None = None
+
+    @staticmethod
+    def _load_chunks(path: Path) -> list[dict]:
+        out = []
+        with path.open("r", encoding="utf-8") as f:
+            for line in f:
+                out.append(json.loads(line))
+        return out
+
+    def _bm25_search(self, query: str, k: int) -> list[ChunkRec]:
+        scores = self.bm25.get_scores(tokenize(query))
+        idx = np.argsort(scores)[::-1][:k]
+        out: list[ChunkRec] = []
+        for i in idx:
+            c = self.chunks[int(i)]
+            out.append(
+                ChunkRec(
+                    c["chunk_id"],
+                    c["source_id"],
+                    c["text"],
+                    float(scores[int(i)]),
+                    "bm25",
+                )
+            )
+        return out
+
+    def _dense_search(self, query: str, k: int) -> list[ChunkRec]:
+        q = self.embedder.encode([query], normalize_embeddings=True)
+        q = np.asarray(q, dtype=np.float32)[0]
+        # cosine similarity because embeddings normalized
+        scores = self.emb @ q
+        idx = np.argsort(scores)[::-1][:k]
+        out: list[ChunkRec] = []
+        for i in idx:
+            c = self.chunks[int(i)]
+            out.append(
+                ChunkRec(
+                    c["chunk_id"],
+                    c["source_id"],
+                    c["text"],
+                    float(scores[int(i)]),
+                    "dense",
+                )
+            )
+        return out
+
+    def _get_reranker(self) -> CrossEncoder:
+        if self._reranker is None:
+            self._reranker = CrossEncoder(SETTINGS.rerank_model)
+        return self._reranker
+
+    def retrieve(
+        self,
+        query: str,
+        use_bm25: bool = True,
+        use_dense: bool = True,
+        use_rerank: bool = False,
+    ) -> list[ChunkRec]:
+        cands: list[ChunkRec] = []
+        if use_bm25:
+            cands.extend(self._bm25_search(query, SETTINGS.top_k_bm25))
+        if use_dense:
+            cands.extend(self._dense_search(query, SETTINGS.top_k_dense))
+
+        # de-dup by chunk_id keeping best score per chunk
+        best: dict[int, ChunkRec] = {}
+        for r in cands:
+            prev = best.get(r.chunk_id)
+            if prev is None or r.score > prev.score:
+                best[r.chunk_id] = r
+        merged = list(best.values())
+        merged.sort(key=lambda x: x.score, reverse=True)
+
+        if use_rerank and merged:
+            reranker = self._get_reranker()
+            top = merged[: SETTINGS.rerank_top_n]
+            pairs = [(query, r.text) for r in top]
+            rr_scores = reranker.predict(pairs)
+            for r, s in zip(top, rr_scores):
+                r.score = float(s)
+                r.why = "rerank"
+            top.sort(key=lambda x: x.score, reverse=True)
+            return top[: SETTINGS.top_k_final]
+
+        return merged[: SETTINGS.top_k_final]

requirements.txt

@@ -0,0 +1,87 @@
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.13.2
+aiosignal==1.4.0
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.12.0
+attrs==25.4.0
+audioop-lts==0.2.2
+brotli==1.2.0
+certifi==2025.11.12
+charset-normalizer==3.4.4
+click==8.3.1
+datasets==4.4.1
+dill==0.4.0
+distro==1.9.0
+fastapi==0.124.4
+ffmpy==1.0.0
+filelock==3.20.0
+frozenlist==1.8.0
+fsspec==2025.10.0
+gradio==6.1.0
+gradio_client==2.0.1
+groovy==0.1.2
+h11==0.16.0
+hf-xet==1.2.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.36.0
+idna==3.11
+Jinja2==3.1.6
+jiter==0.12.0
+joblib==1.5.2
+markdown-it-py==4.0.0
+MarkupSafe==3.0.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.7.0
+multiprocess==0.70.18
+networkx==3.6.1
+numpy==2.3.5
+openai==2.11.0
+orjson==3.11.5
+packaging==25.0
+pandas==2.3.3
+pillow==12.0.0
+propcache==0.4.1
+pyarrow==22.0.0
+pydantic==2.12.4
+pydantic_core==2.41.5
+pydub==0.25.1
+Pygments==2.19.2
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytz==2025.2
+PyYAML==6.0.3
+rank-bm25==0.2.2
+regex==2025.11.3
+requests==2.32.5
+rich==14.2.0
+safehttpx==0.1.7
+safetensors==0.7.0
+scikit-learn==1.8.0
+scipy==1.16.3
+semantic-version==2.10.0
+sentence-transformers==5.2.0
+setuptools==80.9.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+starlette==0.50.0
+sympy==1.14.0
+threadpoolctl==3.6.0
+tokenizers==0.22.1
+tomlkit==0.13.3
+torch==2.9.1
+tqdm==4.67.1
+transformers==4.57.3
+typer==0.20.0
+typer-slim==0.20.0
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+tzdata==2025.3
+urllib3==2.6.2
+uvicorn==0.38.0
+xxhash==3.6.0
+yarl==1.22.0