Document Search × Python── Searching large document collections by meaning, not just matching words

"""Keep a search corpus in step with a drive: fetch only new or edited files, drop deleted ones."""
import asyncio
import json
import os
from pathlib import Path

from azure.identity import DeviceCodeCredential
from msgraph import GraphServiceClient

SCOPES = ["Files.Read.All"]
EXTS = (".docx", ".pptx", ".pdf")
STATE = Path(os.environ.get("INDEX_STATE", "out_index_state.json"))
RAW = Path(os.environ.get("RAW_DIR", "out_raw"))


def plan_changes(items, known):
    """items: delta results as dicts; known: {item id: cTag} from the last run."""
    fetch, drop = [], []
    for it in items:
        if it.get("deleted") is not None:
            if it["id"] in known:
                drop.append(it["id"])
            continue
        if not it.get("file") or not it.get("name", "").lower().endswith(EXTS):
            continue  # folders, images and other types are not indexed
        if known.get(it["id"]) != it.get("cTag"):
            fetch.append(it)  # new file, or its content changed (cTag ignores metadata-only edits)
    return fetch, drop


async def read_delta(builder, link):
    items, page = [], await builder.with_url(link).get()
    while True:
        for i in page.value or []:
            items.append({"id": i.id, "name": i.name or "", "cTag": i.c_tag,
                          "file": i.file is not None, "deleted": i.deleted})
        if not page.odata_next_link:
            return items, page.odata_delta_link
        page = await builder.with_url(page.odata_next_link).get()


async def main():
    drive_id = os.environ["DRIVE_ID"]
    cred = DeviceCodeCredential(client_id=os.environ["CLIENT_ID"], tenant_id=os.environ["TENANT_ID"])
    client = GraphServiceClient(credentials=cred, scopes=SCOPES)
    state = {"files": {}, "deltaLink": None}
    if STATE.exists():
        state = json.loads(STATE.read_text(encoding="utf-8"))
    items_rb = client.drives.by_drive_id(drive_id).items
    link = state["deltaLink"] or "https://graph.microsoft.com/v1.0/drives/%s/root/delta" % drive_id
    items, delta_link = await read_delta(items_rb.by_drive_item_id("root").delta, link)
    fetch, drop = plan_changes(items, state["files"])
    RAW.mkdir(exist_ok=True)
    for it in fetch:
        data = await items_rb.by_drive_item_id(it["id"]).content.get()
        (RAW / (it["id"] + Path(it["name"]).suffix.lower())).write_bytes(data)
        state["files"][it["id"]] = it["cTag"]
    for item_id in drop:
        state["files"].pop(item_id, None)  # also remove this id's chunks from your index
    state["deltaLink"] = delta_link
    STATE.write_text(json.dumps(state, indent=1), encoding="utf-8")
    print("to (re)index:", [it["name"] for it in fetch], "to drop:", drop)


if __name__ == "__main__":
    asyncio.run(main())

"""Extract text from one PDF with pypdf, pdfplumber and PyMuPDF, and compare what each returns."""
import pdfplumber
import pymupdf
from pypdf import PdfReader

PATH = "samples/search_policy.pdf"


def with_pypdf(path):
    return [page.extract_text() or "" for page in PdfReader(path).pages]


def with_pdfplumber(path):
    with pdfplumber.open(path) as pdf:
        return [page.extract_text() or "" for page in pdf.pages]


def with_pymupdf(path):
    with pymupdf.open(path) as doc:
        return [page.get_text() for page in doc]


for name, extract in [("pypdf", with_pypdf), ("pdfplumber", with_pdfplumber), ("pymupdf", with_pymupdf)]:
    pages = extract(PATH)
    chars = sum(len(t.strip()) for t in pages)
    lines = pages[0].strip().splitlines() if pages else []
    print("%-10s pages=%d chars=%d first line=%s" % (name, len(pages), chars, lines[0] if lines else "-"))
    thin = [n for n, t in enumerate(pages, start=1) if len(t.strip()) < 20]
    if thin:
        print("  little or no text on pages", thin, "- possibly scanned; OCR is needed")

"""Chunk documents two ways: by Word headings, and by fixed size with overlap (PDF text)."""
import json

import docx
import pymupdf

SIZE, OVERLAP = 120, 30  # characters; keep chunks well inside the embedding model's input limit


def by_heading(path):
    chunks, heads, buf = [], [], []

    def flush():
        if buf:
            chunks.append({"source": path, "heading": " > ".join(heads), "text": "".join(buf)})
            buf.clear()

    for p in docx.Document(path).paragraphs:
        style = p.style.name if p.style is not None else ""
        if style.startswith("Heading"):
            flush()
            level = int(style.split()[-1]) if style.split()[-1].isdigit() else 1
            heads[:] = heads[:level - 1] + [p.text.strip()]
        elif p.text.strip():
            buf.append(p.text.strip())
    flush()
    return chunks


def fixed_size(text, source, size=SIZE, overlap=OVERLAP):
    text = "".join(line.strip() for line in text.splitlines())  # Japanese: join lines without spaces
    step = size - overlap
    return [{"source": source, "start": i, "text": text[i:i + size]}
            for i in range(0, max(len(text) - overlap, 1), step)]


chunks = by_heading("samples/search_handbook.docx")
with pymupdf.open("samples/search_policy.pdf") as pdf:
    for n, page in enumerate(pdf, start=1):
        chunks += fixed_size(page.get_text(), "search_policy.pdf#page=%d" % n)
with open("out_chunks.jsonl", "w", encoding="utf-8") as f:
    for c in chunks:
        f.write(json.dumps(c, ensure_ascii=False) + "\n")
print("chunks:", len(chunks))
for c in chunks[:3] + chunks[-2:]:
    print("-", c.get("heading") or "%s@%d" % (c["source"], c["start"]), "|", c["text"][:40])

"""Keyword baselines for Japanese text: TF-IDF on character n-grams, and BM25 on character bigrams."""
import json

from rank_bm25 import BM25Okapi
from sklearn.feature_extraction.text import TfidfVectorizer

with open("samples/search_corpus.jsonl", encoding="utf-8") as f:
    docs = [json.loads(line) for line in f]
texts = [d["title"] + " " + d["text"] for d in docs]


def bigrams(s):
    s = "".join(s.split()).lower()  # no word spaces in Japanese: use overlapping 2-character tokens
    return [s[i:i + 2] for i in range(len(s) - 1)]


tfidf = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 3), sublinear_tf=True)
X = tfidf.fit_transform(texts)  # rows are L2-normalised, so a dot product is cosine similarity
bm25 = BM25Okapi([bigrams(t) for t in texts])


def top(scores, k=3):
    order = sorted(range(len(scores)), key=lambda i: -scores[i])[:k]
    return ", ".join("%s %s (%.2f)" % (docs[i]["id"], docs[i]["title"], scores[i]) for i in order)


for q in ["経費の申請の締め日", "家で仕事をする日の届け出"]:
    s_tfidf = (X @ tfidf.transform([q]).T).toarray().ravel()
    s_bm25 = bm25.get_scores(bigrams(q))
    print("Q:", q)
    print("  tfidf:", top(s_tfidf))
    print("  bm25 :", top(s_bm25))

"""Semantic search: embed documents with a multilingual model and search them with a FAISS index."""
import json

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

MODEL = "intfloat/multilingual-e5-small"  # 384-dimensional vectors (see the model card)

with open("samples/search_corpus.jsonl", encoding="utf-8") as f:
    docs = [json.loads(line) for line in f]
model = SentenceTransformer(MODEL)
# e5 models expect a "passage: " prefix on documents and "query: " on queries
vecs = model.encode(["passage: " + d["title"] + " " + d["text"] for d in docs],
                    normalize_embeddings=True, batch_size=16)
vecs = np.asarray(vecs, dtype="float32")
index = faiss.IndexFlatIP(vecs.shape[1])  # inner product of unit vectors = cosine similarity
index.add(vecs)  # vector i <-> docs[i]; keep this id mapping next to the saved index
faiss.write_index(index, "out_corpus.faiss")
print("dim:", vecs.shape[1], "vectors:", index.ntotal)

for q in ["家で仕事をする日の届け出", "ログインの合言葉を忘れてしまった"]:
    qv = np.asarray(model.encode(["query: " + q], normalize_embeddings=True), dtype="float32")
    scores, ids = index.search(qv, 3)
    print("Q:", q)
    for s, i in zip(scores[0], ids[0]):
        print("  %.3f %s %s" % (s, docs[i]["id"], docs[i]["title"]))

"""Hybrid search: fuse a BM25 ranking and an embedding ranking with Reciprocal Rank Fusion (RRF)."""
import json

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

K_RRF = 60  # the constant k in 1/(rank + k)

with open("samples/search_corpus.jsonl", encoding="utf-8") as f:
    docs = [json.loads(line) for line in f]
texts = [d["title"] + " " + d["text"] for d in docs]


def bigrams(s):
    s = "".join(s.split()).lower()
    return [s[i:i + 2] for i in range(len(s) - 1)]


bm25 = BM25Okapi([bigrams(t) for t in texts])
model = SentenceTransformer("intfloat/multilingual-e5-small")
emb = model.encode(["passage: " + t for t in texts], normalize_embeddings=True)


def rrf(rankings, k=K_RRF, weights=None):
    fused = {}
    for n, ranking in enumerate(rankings):
        w = 1.0 if weights is None else weights[n]
        for rank, doc in enumerate(ranking, start=1):
            fused[doc] = fused.get(doc, 0.0) + w / (rank + k)
    return sorted(fused, key=fused.get, reverse=True)


def search(q, depth=20):
    lexical = [int(i) for i in np.argsort(-bm25.get_scores(bigrams(q)))[:depth]]
    qv = model.encode(["query: " + q], normalize_embeddings=True)[0]
    semantic = [int(i) for i in np.argsort(-(emb @ qv))[:depth]]
    return {"bm25": lexical, "embed": semantic, "hybrid": rrf([lexical, semantic])}


for q in ["VPNにつながらない", "家で仕事をする日の届け出", "高い機械を買うときに必要な承認"]:
    print("Q:", q)
    for method, ranking in search(q).items():
        print("  %-6s %s" % (method, " ".join(docs[i]["id"] for i in ranking[:3])))

"""Rerank the top candidates from a fast first-stage search with a cross-encoder."""
import json

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import CrossEncoder

RERANKER = "cl-nagoya/ruri-v3-reranker-310m"  # Japanese reranker, Apache 2.0 (model card)
DEPTH = 10  # how many first-stage candidates the reranker reads

with open("samples/search_corpus.jsonl", encoding="utf-8") as f:
    docs = [json.loads(line) for line in f]
texts = [d["title"] + " " + d["text"] for d in docs]


def bigrams(s):
    s = "".join(s.split()).lower()
    return [s[i:i + 2] for i in range(len(s) - 1)]


bm25 = BM25Okapi([bigrams(t) for t in texts])
reranker = CrossEncoder(RERANKER)

for q in ["立て替えたお金を返してもらう手続きの期限", "高い機械を買うときに必要な承認"]:
    cand = [int(i) for i in np.argsort(-bm25.get_scores(bigrams(q)))[:DEPTH]]
    ranked = reranker.rank(q, [texts[i] for i in cand], top_k=3)
    print("Q:", q)
    print("  bm25  :", " ".join(docs[i]["id"] for i in cand[:3]))
    print("  rerank:", " ".join("%s(%.2f)" % (docs[cand[r["corpus_id"]]]["id"], r["score"]) for r in ranked))

"""Get grounding extracts from SharePoint through the Microsoft 365 Copilot Retrieval API (read-only)."""
import asyncio
import json
import os

from azure.identity import DeviceCodeCredential
from kiota_abstractions.method import Method
from kiota_abstractions.request_information import RequestInformation
from msgraph import GraphServiceClient

SCOPES = ["Files.Read.All", "Sites.Read.All"]
URL = "https://graph.microsoft.com/v1.0/copilot/retrieval"


def build_body(query, site_url=None, max_results=10):
    body = {"queryString": query, "dataSource": "sharePoint",
            "resourceMetadata": ["title", "author"], "maximumNumberOfResults": max_results}
    if site_url:
        body["filterExpression"] = 'path:"%s"' % site_url  # KQL: only this site
    return body


def to_rows(resp):
    rows = []
    for hit in resp.get("retrievalHits", []):
        meta = hit.get("resourceMetadata") or {}
        label = (hit.get("sensitivityLabel") or {}).get("displayName", "")
        for ex in hit.get("extracts", []):
            rows.append({"title": meta.get("title", ""), "url": hit.get("webUrl", ""), "label": label,
                         "score": ex.get("relevanceScore"), "text": ex.get("text", "")})
    return rows


async def main():
    cred = DeviceCodeCredential(client_id=os.environ["CLIENT_ID"], tenant_id=os.environ["TENANT_ID"])
    client = GraphServiceClient(credentials=cred, scopes=SCOPES)
    body = build_body(os.environ.get("QUERY", "在宅勤務を申請するときの期限と手順を知りたい"),
                      os.environ.get("SITE_URL"))
    req = RequestInformation(Method.POST)  # POST carries the query; nothing in Microsoft 365 is changed
    req.url = URL
    req.set_stream_content(json.dumps(body, ensure_ascii=False).encode("utf-8"), "application/json")
    raw = await client.request_adapter.send_primitive_async(req, "bytes", None)
    for r in to_rows(json.loads(raw)):
        print("%s | %s | %s" % (r["title"], r["label"] or "-", r["text"][:60]))


if __name__ == "__main__":
    asyncio.run(main())

"""Hybrid (keyword + vector) query against an Azure AI Search index, signed in with Microsoft Entra ID."""
import os

from azure.identity import DefaultAzureCredential
from azure.search.documents import SearchClient
from azure.search.documents.models import VectorizableTextQuery

ENDPOINT = os.environ.get("SEARCH_ENDPOINT", "https://example.search.windows.net")
INDEX = os.environ.get("SEARCH_INDEX", "docs-chunks")


def hybrid(client, text, k=5):
    # the index's vector field needs a vectorizer so the service can embed the query text
    vq = VectorizableTextQuery(text=text, k_nearest_neighbors=50, fields="contentVector")
    return client.search(search_text=text, vector_queries=[vq], select=["title", "url", "chunk"], top=k)


def main():
    client = SearchClient(ENDPOINT, INDEX, DefaultAzureCredential())  # needs "Search Index Data Reader"
    for r in hybrid(client, "在宅勤務を申請するときの期限"):
        print("%.4f | %s | %s" % (r["@search.score"], r["title"], r["url"]))


if __name__ == "__main__":
    main()

"""End to end: search one corpus with several methods and score each with recall@k and MRR."""
import csv
import json
import os

import numpy as np
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction.text import TfidfVectorizer

KS = (1, 3, 5)
DEPTH = 10  # a relevant document ranked below this counts as not found (MRR@10)

with open("samples/search_corpus.jsonl", encoding="utf-8") as f:
    docs = [json.loads(line) for line in f]
with open("samples/search_queries.json", encoding="utf-8") as f:
    queries = json.load(f)  # [{"id", "query", "relevant": [doc ids]}]
texts = [d["title"] + " " + d["text"] for d in docs]
ids = [d["id"] for d in docs]


def bigrams(s):
    s = "".join(s.split()).lower()
    return [s[i:i + 2] for i in range(len(s) - 1)]


def ranking(scores):
    return [ids[i] for i in np.argsort(-np.asarray(scores))[:DEPTH]]


def rrf(rankings, k=60):
    fused = {}
    for r in rankings:
        for rank, d in enumerate(r, start=1):
            fused[d] = fused.get(d, 0.0) + 1.0 / (rank + k)
    return sorted(fused, key=fused.get, reverse=True)[:DEPTH]


def first_hit(got, relevant):
    return next((n for n, d in enumerate(got, start=1) if d in relevant), None)


def evaluate(results):
    """results: {query id: ranked doc ids} -> mean recall@k for each k in KS, and MRR@DEPTH."""
    out = {"recall@%d" % k: 0.0 for k in KS}
    out["MRR@%d" % DEPTH] = 0.0
    for q in queries:
        got, rel = results[q["id"]], set(q["relevant"])
        for k in KS:
            out["recall@%d" % k] += len(rel & set(got[:k])) / len(rel) / len(queries)
        r = first_hit(got, rel)
        out["MRR@%d" % DEPTH] += (1.0 / r if r else 0.0) / len(queries)
    return out


methods = {}
tfidf = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 3), sublinear_tf=True)
X = tfidf.fit_transform(texts)
methods["tfidf"] = {q["id"]: ranking((X @ tfidf.transform([q["query"]]).T).toarray().ravel()) for q in queries}
bm25 = BM25Okapi([bigrams(t) for t in texts])
methods["bm25"] = {q["id"]: ranking(bm25.get_scores(bigrams(q["query"]))) for q in queries}

try:  # the semantic steps need the model files; set HF_HUB_OFFLINE=1 to stay offline
    import faiss
    from sentence_transformers import CrossEncoder, SentenceTransformer
    model = SentenceTransformer(os.environ.get("EMBED_MODEL", "intfloat/multilingual-e5-small"))
except Exception as e:
    model = None
    print("embedding skipped:", type(e).__name__, e)
if model is not None:
    emb = model.encode(["passage: " + t for t in texts], normalize_embeddings=True)
    index = faiss.IndexFlatIP(emb.shape[1])
    index.add(np.asarray(emb, dtype="float32"))
    qv = model.encode(["query: " + q["query"] for q in queries], normalize_embeddings=True)
    _, hits = index.search(np.asarray(qv, dtype="float32"), DEPTH)
    methods["embed"] = {q["id"]: [ids[i] for i in hits[n]] for n, q in enumerate(queries)}
    methods["hybrid"] = {q["id"]: rrf([methods["bm25"][q["id"]], methods["embed"][q["id"]]]) for q in queries}
    if os.environ.get("RERANK", "0") == "1":  # slow on CPU: opt in with RERANK=1
        try:
            ce = CrossEncoder(os.environ.get("RERANK_MODEL", "cl-nagoya/ruri-v3-reranker-310m"))
            by_id = dict(zip(ids, texts))
            reranked = {}
            for q in queries:
                cand = methods["hybrid"][q["id"]]
                order = ce.rank(q["query"], [by_id[d] for d in cand])
                reranked[q["id"]] = [cand[r["corpus_id"]] for r in order]
            methods["hybrid+rerank"] = reranked
        except Exception as e:
            print("rerank skipped:", type(e).__name__, e)

names = list(methods)
scores = {n: evaluate(methods[n]) for n in names}
cols = list(scores[names[0]])
print("%-14s" % "method" + "".join("%10s" % c for c in cols))
for n in names:
    print("%-14s" % n + "".join("%10.3f" % scores[n][c] for c in cols))
with open("out_search_eval.csv", "w", newline="", encoding="utf-8") as f:
    w = csv.writer(f)
    w.writerow(["method"] + cols)
    w.writerows([[n] + [round(scores[n][c], 4) for c in cols] for n in names])
print("\nrank of the first relevant document per query (- = not in top %d)" % DEPTH)
for q in queries:
    cells = [str(first_hit(methods[n][q["id"]], set(q["relevant"])) or "-") for n in names]
    print(q["id"], " ".join("%s=%s" % (n, c) for n, c in zip(names, cells)), "|", q["query"])