RAG Part 3: Vector Store

import sys, os, glob, ctypes, subprocess, importlib.util
from importlib.metadata import version, PackageNotFoundError

# ---- config ----
CUDA_WHL = "https://abetlen.github.io/llama-cpp-python/whl/cu121"  # switch to cu122 if your DC CUDA is 12.2
LLAMA_VER = "0.3.16"

BASE_PKGS = ["datasets", "onnxruntime-gpu", "transformers", "huggingface_hub"]
CUDA_PKGS = ["nvidia-cuda-runtime-cu12", "nvidia-cublas-cu12"]

# Keep HF quieter / more reliable in demos
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "120"

def have_module(modname: str) -> bool:
    return importlib.util.find_spec(modname) is not None

def have_pkg(pkgname: str, want: str | None = None) -> bool:
    try:
        v = version(pkgname)
        return (want is None) or (v == want)
    except PackageNotFoundError:
        return False

def pip_install_quiet(args: list[str]):
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-U", "--no-cache-dir", "-q", "--break-system-packages", *args],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

def pip_uninstall_quiet(pkg: str):
    subprocess.run(
        [sys.executable, "-m", "pip", "uninstall", "-y", pkg],
        check=False,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

# ---- 1) install only if missing ----
missing = [p for p in BASE_PKGS if not have_pkg(p)]
if missing:
    pip_install_quiet(missing)

missing_cuda = [p for p in CUDA_PKGS if not have_pkg(p)]
if missing_cuda:
    pip_install_quiet(missing_cuda)

# llama-cpp-python: ensure correct CUDA wheel + version
need_llama = (not have_pkg("llama-cpp-python", LLAMA_VER))
if need_llama:
    pip_uninstall_quiet("llama-cpp-python")
    pip_install_quiet(["llama-cpp-python", "--extra-index-url", CUDA_WHL])

# ---- 2) preload CUDA .so so import works after kernel restarts ----
cudart = glob.glob("/srv/conda/envs/notebook/lib/python3.12/site-packages/nvidia/**/libcudart.so.12*", recursive=True)
cublas = glob.glob("/srv/conda/envs/notebook/lib/python3.12/site-packages/nvidia/**/libcublas.so.12*", recursive=True)
if cudart and cublas:
    ctypes.CDLL(cudart[0])
    ctypes.CDLL(cublas[0])

# ---- 3) quick sanity line ----
print("Setup OK (deps present). If this was the first run today, restart kernel once, then run again.")

import os, sys
from contextlib import redirect_stderr
from llama_cpp import Llama

with open(os.devnull, "w") as fnull, redirect_stderr(fnull):
    model = Llama(
        model_path="/home/jovyan/shared/qwen2-1_5b-instruct-q4_0.gguf",
        n_ctx=2048,
        n_gpu_layers=28,
        verbose=True,
    )

# --- Cell 3: runtime imports + embedding model (no llama install here) ---

import os, time, glob, ctypes
import numpy as np
import onnxruntime as ort
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

# HF settings
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "120"

# Preload CUDA runtime libs so llama-cpp-python can import after kernel restarts
cudart = glob.glob("/srv/conda/envs/notebook/lib/python3.12/site-packages/nvidia/**/libcudart.so.12*", recursive=True)
cublas = glob.glob("/srv/conda/envs/notebook/lib/python3.12/site-packages/nvidia/**/libcublas.so.12*", recursive=True)
if cudart and cublas:
    ctypes.CDLL(cudart[0])
    ctypes.CDLL(cublas[0])

from llama_cpp import Llama  # now safe

# Download ONNX embedding model
onnx_path = hf_hub_download(
    repo_id="sentence-transformers/all-MiniLM-L6-v2",
    filename="onnx/model_O4.onnx",
    force_download=True
)

# Prefer CUDA EP if available (priority order matters)  :contentReference[oaicite:0]{index=0}
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in ort.get_available_providers() else ["CPUExecutionProvider"]
onnx_session = ort.InferenceSession(onnx_path, providers=providers)
active_provider = onnx_session.get_providers()[0]
print(f"ONNX active provider: {active_provider}")

tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

def encode(texts):
    if isinstance(texts, str):
        texts = [texts]
    encoded = tokenizer(texts, padding=True, truncation=True, max_length=256, return_tensors="np")
    outputs = onnx_session.run(None, {
        "input_ids": encoded["input_ids"],
        "attention_mask": encoded["attention_mask"],
        "token_type_ids": encoded["token_type_ids"],
    })
    token_embs = outputs[0]
    mask = np.expand_dims(encoded["attention_mask"], axis=-1).astype(float)
    emb = np.sum(token_embs * mask, axis=1) / np.maximum(mask.sum(axis=1), 1e-9)
    emb = emb / np.linalg.norm(emb, axis=1, keepdims=True)
    return emb

# Quick warmup
t0 = time.perf_counter()
_ = encode(["warm up the embedding model"])
t1 = time.perf_counter()
print(f"Setup complete! (embedding warmup: {t1-t0:.3f}s on {active_provider})")

# Load the dataset from Hugging Face
ds = load_dataset("meetara-lab/vectorstore-economics", split="train")

print(f"Loaded {len(ds)} raw chunks from the dataset")
print(f"Columns: {ds.column_names}")
print()

# Let's look at what a raw chunk looks like
print("Example chunk (raw):")
print("=" * 50)
print(ds[100]["document"][:400])

# Filter out boilerplate — copyright notices, table of contents, contributor lists, etc.
noise_markers = [
    "creative commons", "openstax", "all rights reserved", "rice university",
    "table of contents", "contributor", "peer reviewer", "senior contributing",
    "isbn", "trademarks", "access for free", "citation information",
    "acknowledgment", "bring it home", "link it up", "clear it up", "work it out"
]

chunks = []
chunk_embeddings_list = []
for item in ds:
    text = item["document"]
    text_lower = text.lower()

    # Skip boilerplate and chunks with very few unique words
    is_noise = any(marker in text_lower for marker in noise_markers)
    too_repetitive = len(set(text.split())) < 30

    if not is_noise and not too_repetitive:
        chunks.append(text)
        chunk_embeddings_list.append(item["embedding"])

print(f"After filtering: {len(chunks)} usable textbook passages (removed {len(ds) - len(chunks)} noisy chunks)")
print()

# Use a subset to keep things fast
MAX_CHUNKS = 500
chunks = chunks[:MAX_CHUNKS]
chunk_embeddings = np.array(chunk_embeddings_list[:MAX_CHUNKS], dtype=np.float32)

# Normalize (in case the stored embeddings aren't already)
norms = np.linalg.norm(chunk_embeddings, axis=1, keepdims=True)
chunk_embeddings = chunk_embeddings / np.maximum(norms, 1e-9)

print(f"Using first {len(chunks)} chunks with pre-computed embeddings")
print(f"Vector store shape: {chunk_embeddings.shape} (no embedding step needed!)")

# Let's peek at what our textbook passages actually look like
print("Sample textbook passages:")
print("=" * 60)
for i in [0, 50, 150, 300, 450]:
    if i < len(chunks):
        preview = chunks[i][:200].replace("\n", " ")
        print(f"\nChunk {i}:")
        print(f"  {preview}...")

def search_textbook(question, n_results=3):
    """Search the vector store for the most relevant textbook passages."""
    # Embed the question
    q_embedding = encode([question])

    # Compare to all passages (dot product = cosine similarity for normalized vectors)
    similarities = (q_embedding @ chunk_embeddings.T)[0]

    # Get the top matches
    top_indices = np.argsort(similarities)[::-1][:n_results]

    results = []
    for idx in top_indices:
        results.append({
            "index": int(idx),
            "text": chunks[idx],
            "similarity": float(similarities[idx])
        })
    return results

print("Search function ready!")

# Ask a study question
question = "What determines the price of a product in a market?"

results = search_textbook(question, n_results=3)

print(f'Study question: "{question}"')
print()
print("Top 3 textbook passages found:")
print("=" * 60)
for i, r in enumerate(results):
    print(f"\n--- Result #{i+1} (chunk {r['index']}, similarity: {r['similarity']:.3f}) ---")
    print(r["text"][:300])
    if len(r["text"]) > 300:
        print("...")

# Visualize which chunks were retrieved out of all 500
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
retrieved_indices = set(r["index"] for r in results)

fig, ax = plt.subplots(figsize=(14, 3))

bar_colors = ["#2ecc71" if i in retrieved_indices else "#dfe6e9" for i in range(len(chunks))]

ax.bar(range(len(chunks)), [1] * len(chunks), color=bar_colors, edgecolor="none", width=1.0)
ax.set_title(f'Which passages were retrieved for: "{question}"', fontsize=13)
ax.set_xlabel("Passage index (position in textbook)")
ax.set_yticks([])

legend_elements = [
    Patch(facecolor="#2ecc71", edgecolor="black", label="Retrieved"),
    Patch(facecolor="#dfe6e9", edgecolor="black", label="Not retrieved")
]
ax.legend(handles=legend_elements, loc="upper right")

plt.tight_layout()
plt.show()

print("Each bar is one textbook passage. The green ones are what the search found as most relevant.")
print("Notice they can come from anywhere in the textbook — the search is by meaning, not by position.")

def ask_with_rag(question, n_chunks=3):
    """Retrieve textbook passages, then generate an answer."""
    results = search_textbook(question, n_results=n_chunks)
    retrieved_texts = [r["text"] for r in results]

    context = "\n\n".join(retrieved_texts)
    prompt = f"""Answer the question in 1 sentences using the textbook passages below.
Include specific details and numbers from the passages.

Textbook passages:
{context}

Question: {question}

Answer (1-2 sentences):"""
    model._ctx.reset_timings()
    response = model.create_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=150
    )
    model._ctx.print_timings()
    return response["choices"][0]["message"]["content"], results


def ask_without_rag(question):
    """Ask the model directly — no textbook, just its own knowledge."""
    prompt = f"Answer in 1 sentences: {question}"
    model._ctx.reset_timings()
    response = model.create_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=50
    )
    model._ctx.print_timings()
    return response["choices"][0]["message"]["content"]

print("Functions defined!")

# Pointed question: demands a specific number the model can't guess
test_question = "What was the world price of coffee per pound in 1993?"

answer_no_rag = ask_without_rag(test_question)
answer_rag, rag_results = ask_with_rag(test_question)

print("=" * 60)
print(f'Question: "{test_question}"')
print("=" * 60)
print()
print("WITHOUT RAG (closed-book — model guesses):")
print("-" * 60)
print(answer_no_rag)
print()
print("WITH RAG (open-book — model reads the textbook):")
print("-" * 60)
print(answer_rag)

# Fact check
print()
print("FACT CHECK — the textbook says:")
print("-" * 60)
for r in rag_results:
    if "1993" in r["text"] or "50 cents" in r["text"]:
        print(f"  {r['text'][:400].replace(chr(10), ' ')}")
        break
else:
    print(f"  {rag_results[0]['text'][:400].replace(chr(10), ' ')}")

# Another pointed question — different fact, same pattern
hard_question = "In what year did a major frost hit the Brazilian coffee crop?"

answer_no_rag_2 = ask_without_rag(hard_question)
answer_rag_2, rag_results_2 = ask_with_rag(hard_question)

print("=" * 60)
print(f'Question: "{hard_question}"')
print("=" * 60)
print()
print("WITHOUT RAG (closed-book):")
print("-" * 60)
print(answer_no_rag_2)
print()
print("WITH RAG (open-book):")
print("-" * 60)
print(answer_rag_2)
import re

print()
print("FACT CHECK — the textbook says:")
print("-" * 60)

found = False

for r in rag_results_2:
    idx = r["index"]

    # Join neighboring chunks so we don't get a sentence cut in half
    start_i = max(0, idx - 1)
    end_i = min(len(chunks), idx + 2)
    joined = " ".join(chunks[start_i:end_i])

    # Clean whitespace
    text_flat = re.sub(r"\s+", " ", joined).strip()

    # Find the key phrase (prefer frost + 1994)
    anchor = None
    if "frost" in text_flat.lower() and "1994" in text_flat:
        m = re.search(r"frost", text_flat, flags=re.IGNORECASE)
        if m:
            anchor = m.start()

    if anchor is not None:
        # Print a small passage around the match
        window = 140
        s = max(0, anchor - window)
        e = min(len(text_flat), anchor + window)

        snippet = text_flat[s:e]
        if s > 0:
            snippet = "..." + snippet
        if e < len(text_flat):
            snippet = snippet + "..."

        print(snippet)
        found = True
        break

if not found:
    print("(Couldn’t find a nearby frost/1994 passage in retrieved chunks.)")

context_question = "What happens when the government sets a price floor?"

answer_1, results_1 = ask_with_rag(context_question, n_chunks=1)
answer_3, results_3 = ask_with_rag(context_question, n_chunks=3)
print("=" * 60)
print(f'Question: "{context_question}"')
print("=" * 60)
print()
print("WITH 1 PASSAGE:")
print(f"  Retrieved: chunk {results_1[0]['index']} (similarity {results_1[0]['similarity']:.3f})")
print(f"  Answer: {answer_1}")
print()
print("WITH 3 PASSAGES:")
for r in results_3:
    print(f"  Retrieved: chunk {r['index']} (similarity {r['similarity']:.3f})")
print(f"  Answer: {answer_3}")

# Side-by-side visualization
fig, axes = plt.subplots(1, 2, figsize=(14, 3))

for ax, res, label in [(axes[0], results_1, "1 Passage"), (axes[1], results_3, "3 Passages")]:
    id_set = set(r["index"] for r in res)
    bar_colors = ["#2ecc71" if i in id_set else "#dfe6e9" for i in range(len(chunks))]
    ax.bar(range(len(chunks)), [1] * len(chunks), color=bar_colors, edgecolor="none", width=1.0)
    ax.set_title(label, fontsize=14, fontweight="bold")
    ax.set_xlabel("Passage index")
    ax.set_yticks([])

legend_elements = [
    Patch(facecolor="#2ecc71", edgecolor="black", label="Retrieved"),
    Patch(facecolor="#dfe6e9", edgecolor="black", label="Not retrieved")
]
axes[1].legend(handles=legend_elements, loc="upper right")

plt.suptitle(f'1 vs 3 passages: "{context_question}"', fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 5, figsize=(20, 4))

panels = [
    ("1. Textbook", f"OpenStax\nMicroeconomics 3e\n{len(ds)} raw passages", "#AED6F1", "#2980B9"),
    ("2. Filter", f"{len(chunks)} clean chunks\nremoved boilerplate", "#A9DFBF", "#27AE60"),
    ("3. Embed", f"{chunk_embeddings.shape[0]} vectors\n{chunk_embeddings.shape[1]} dimensions each", "#F9E79F", "#F39C12"),
    ("4. Retrieve", f"Top-k passages\nby cosine similarity", "#F5CBA7", "#E67E22"),
    ("5. Answer", f"Qwen2 1.5B\ncontext + question", "#D7BDE2", "#8E44AD"),
]

for ax, (title, text, facecolor, edgecolor) in zip(axes, panels):
    ax.text(0.5, 0.5, text, ha="center", va="center", fontsize=12,
            bbox=dict(boxstyle="round,pad=0.8", facecolor=facecolor,
                      edgecolor=edgecolor, linewidth=2),
            transform=ax.transAxes)
    ax.set_title(title, fontsize=13, fontweight="bold")
    ax.axis("off")

for i in range(4):
    fig.text(0.19 + i * 0.19, 0.5, "\u2192", fontsize=30, ha="center", va="center")

plt.suptitle("The RAG Pipeline: Textbook to Study Buddy", fontsize=16, fontweight="bold", y=1.05)
plt.tight_layout()
plt.show()

# CHANGE THIS QUESTION!
my_question = "How much did the price of coffee fall to per pound by 2001?"

my_no_rag = ask_without_rag(my_question)
my_rag, my_results = ask_with_rag(my_question, n_chunks=3)

print("=" * 60)
print(f'Your question: "{my_question}"')
print("=" * 60)
print()
print("WITHOUT textbook:")
print("-" * 60)
print(my_no_rag)
print()
print("WITH textbook (RAG):")
print("-" * 60)
print(my_rag)
print()
print("FACT CHECK — the textbook says:")
print("-" * 60)
print(f"  {my_results[0]['text'][:400].replace(chr(10), ' ')}")

# Visualize retrieval
fig, ax = plt.subplots(figsize=(14, 3))
retrieved_set = set(r["index"] for r in my_results)
bar_colors = ["#2ecc71" if i in retrieved_set else "#dfe6e9" for i in range(len(chunks))]
ax.bar(range(len(chunks)), [1] * len(chunks), color=bar_colors, edgecolor="none", width=1.0)
ax.set_title(f'Retrieved passages for: "{my_question}"', fontsize=13)
ax.set_xlabel("Passage index")
ax.set_yticks([])
plt.tight_layout()
plt.show()