Chapter 10: The Capstone

import numpy as np
import re
import math

# The full pipeline library from Chapters 2-9, condensed.
training_corpus = [
    "Employees may carry over unused vacation days into the next calendar year.",
    "Staff can roll over unused holiday days into next year.",
    "Workers get paid time off and vacation days every year.",
    "Most employees take a vacation or holiday once a year.",
    "Unused vacation time can roll over to next year for many employees.",
    "Holiday leave and vacation leave are treated the same way at most companies.",
    "Some employees prefer to use their holiday days instead of saving them.",
    "Vacation days and holiday days both count as paid time off.",
    "Expense reports must be submitted within days of the purchase.",
    "Employees submit an expense report after every business purchase.",
    "The finance team reviews each expense report and each purchase receipt.",
    "A purchase over the limit needs approval before the expense is reimbursed.",
    "Expense claims and purchase receipts are reviewed together by finance.",
    "Remote employees receive a home office stipend to set up their workspace.",
    "The company gives remote workers a stipend for home office equipment.",
    "Employees working remote from home can request an office stipend.",
    "A stipend helps remote staff pay for office furniture at home.",
    "The office stipend is a one time payment for remote employees.",
    "Vacation policy, expense policy, and remote work policy are all in the handbook.",
    "Every employee should read the handbook for vacation and expense rules.",
    "Employees who are sick should notify their manager as soon as possible.",
    "Sick leave is paid time off for employees who are too unwell to work.",
    "A doctor's note may be required after three consecutive sick days.",
    "Unused sick leave does not roll over into the next year.",
    "Employees should stay home and rest when they are sick.",
    "Standard working hours are nine to five, Monday through Friday.",
    "Employees may request a flexible schedule with manager approval.",
    "Overtime must be approved in advance by a manager.",
    "Most staff work standard business hours unless a flexible schedule is approved.",
    "Employees are expected to treat colleagues and customers with respect.",
    "Harassment or discrimination of any kind will not be tolerated.",
    "Violations of the code of conduct may result in disciplinary action.",
    "Every employee must follow the code of conduct and treat others with respect.",
    "Company laptops must be encrypted and password protected.",
    "Employees should not install unauthorized software on company devices.",
    "Report a lost or stolen laptop to IT immediately.",
    "IT equipment must be returned when an employee leaves the company.",
    "The IT team manages laptops, software, and device security for the company.",
    "Contact the helpdesk if you have a technical problem with your computer.",
    "The support team can help reset passwords and fix login issues.",
    "Call the helpdesk for any urgent technical problem during business hours.",
]

def tokenize(text):
    return re.findall(r"[a-z']+", text.lower())

tokens_per_sentence = [tokenize(s) for s in training_corpus]
vocab = sorted(set(w for sent in tokens_per_sentence for w in sent))
word_index = {w: i for i, w in enumerate(vocab)}
V = len(vocab)
window = 4
cooc = np.zeros((V, V))
for sent in tokens_per_sentence:
    for i, w in enumerate(sent):
        for j in range(max(0, i - window), min(len(sent), i + window + 1)):
            if i != j:
                cooc[word_index[w], word_index[sent[j]]] += 1
X = np.log1p(cooc)
U, S, Vt = np.linalg.svd(X, full_matrices=False)
DIM = 25
word_vectors = U[:, :DIM] * S[:DIM]

def embedding(word):
    return word_vectors[word_index[word]] if word in word_index else None

def embed_text(text):
    words = [w for w in tokenize(text) if w in word_index]
    if not words:
        return np.zeros(DIM)
    return np.mean([embedding(w) for w in words], axis=0)

def cosine_similarity(a, b):
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / norm) if norm else 0.0

handbook_document = """Vacation Policy
Employees may carry over up to 5 unused vacation days into the next calendar year. Time off should be requested from a manager at least two weeks in advance when possible.

Sick Leave Policy
Employees who are too unwell to work should notify their manager as soon as possible. Sick leave is paid time off and does not roll over into the next year. A doctor's note may be required after three consecutive sick days.

Expense Policy
Expense reports must be submitted within 30 days of the purchase date. The finance team reviews each expense report and matching purchase receipt before reimbursement.

Remote Work Policy
Remote employees are provided a one-time home office stipend of 500 euros. The stipend covers office furniture and basic equipment needed to work from home.

Working Hours
Standard working hours are nine to five, Monday through Friday. Employees may request a flexible schedule with manager approval, and any overtime must be approved in advance.

Code of Conduct
Employees are expected to treat colleagues and customers with respect. Harassment or discrimination of any kind will not be tolerated, and violations may result in disciplinary action.

IT Support
For urgent technical problems, contact the helpdesk at extension 4477. For non-urgent requests, submit a ticket through the internal support portal."""

section_chunks = [p.strip() for p in handbook_document.split("\n\n")]

def section_title(chunk):
    return chunk.splitlines()[0]

def word_overlap_score(question, candidate):
    return len(set(question.lower().split()) & set(candidate.lower().split()))

def bi_encoder_score(question, candidate):
    return cosine_similarity(embed_text(question), embed_text(candidate))

class BM25:
    def __init__(self, documents, k1=1.5, b=0.75):
        self.k1, self.b = k1, b
        self.doc_tokens = [re.findall(r"[a-z0-9']+", d.lower()) for d in documents]
        self.doc_lens = [len(t) for t in self.doc_tokens]
        self.avgdl = sum(self.doc_lens) / len(self.doc_lens)
        self.N = len(documents)
        doc_freq = {}
        for toks in self.doc_tokens:
            for term in set(toks):
                doc_freq[term] = doc_freq.get(term, 0) + 1
        self.idf = {t: math.log((self.N - n + 0.5) / (n + 0.5) + 1) for t, n in doc_freq.items()}

    def score(self, query, doc_index):
        term_freq = {}
        for t in self.doc_tokens[doc_index]:
            term_freq[t] = term_freq.get(t, 0) + 1
        dl = self.doc_lens[doc_index]
        total = 0.0
        for term in re.findall(r"[a-z0-9']+", query.lower()):
            if term not in self.idf:
                continue
            f = term_freq.get(term, 0)
            total += self.idf[term] * (f * (self.k1 + 1)) / (f + self.k1 * (1 - self.b + self.b * dl / self.avgdl))
        return total

bm25 = BM25(section_chunks)

def rrf_rank(question, chunks, rrf_k=60):
    embed_ranked = sorted(range(len(chunks)), key=lambda i: bi_encoder_score(question, chunks[i]), reverse=True)
    bm25_ranked = sorted(range(len(chunks)), key=lambda i: bm25.score(question, i), reverse=True)
    scores = {}
    for rank, i in enumerate(embed_ranked, start=1):
        scores[i] = scores.get(i, 0) + 1 / (rrf_k + rank)
    for rank, i in enumerate(bm25_ranked, start=1):
        scores[i] = scores.get(i, 0) + 1 / (rrf_k + rank)
    return sorted(scores.items(), key=lambda kv: kv[1], reverse=True)

def late_interaction_score(question, candidate):
    q_words = [w for w in tokenize(question) if w in word_index]
    c_words = [w for w in tokenize(candidate) if w in word_index]
    if not q_words or not c_words:
        return 0.0
    candidate_vectors = [embedding(w) for w in c_words]
    return sum(max(cosine_similarity(embedding(qw), cv) for cv in candidate_vectors) for qw in q_words) / len(q_words)

def get_top1(question, method, chunks):
    if method == "word_overlap":
        return max(chunks, key=lambda c: word_overlap_score(question, c))
    if method == "embeddings":
        return max(chunks, key=lambda c: bi_encoder_score(question, c))
    if method == "hybrid":
        ranked = rrf_rank(question, chunks)
        return chunks[ranked[0][0]]
    if method == "hybrid_rerank":
        ranked = rrf_rank(question, chunks)
        shortlist = [chunks[i] for i, _ in ranked[:3]]
        return max(shortlist, key=lambda c: late_interaction_score(question, c))
    raise ValueError(f"unknown method: {method}")

print("Pipeline library ready.")