File size: 2,885 Bytes
a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e 3637c38 a43117e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """Manual retrieval-augmented generation for a document Q&A lesson.
Both the embedding model and the generation model run on CPU. The generation
model is loaded only when a user asks for an actual generated answer.
"""
import numpy as np
import torch
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GENERATION_MODEL = "google/flan-t5-small"
MIN_SCORE = 0.30 # Teaching heuristic, NOT calibrated confidence.
NO_ANSWER = "I cannot answer this from the supplied passages."
DATASET_ID = "rajpurkar/squad"
DATASET_REVISION = "7b6d24c440a36b6815f21b70d25016731768db1f"
corpus = []
source_questions = []
embedder = None
passage_embeddings = None
tokenizer = None
generator = None
def load_passage_sample(embedding_tokenizer):
"""Stream a reproducible sample of 30 distinct SQuAD contexts."""
rows = load_dataset(
DATASET_ID,
split="train",
revision=DATASET_REVISION,
streaming=True,
)
passages, questions = [], []
seen_contexts = set()
title_counts = {}
for row in rows.take(6000):
title = row["title"].replace("_", " ")
context = row["context"].strip()
if (
context in seen_contexts
or title_counts.get(title, 0) >= 3
or len(title) > 35
or title == "Genocide"
):
continue
token_count = len(
embedding_tokenizer(
title + ". " + context,
truncation=False,
verbose=False,
)["input_ids"]
)
if token_count > 200:
continue
passage_id = f"P{len(passages) + 1:02d}"
passages.append(
{
"id": passage_id,
"title": title,
"text": context,
"source": "SQuAD train / Wikipedia: " + title,
"squad_id": row["id"],
}
)
questions.append(
{
"query": row["question"],
"expected_ids": [passage_id],
"squad_id": row["id"],
"answer_spans": row["answers"]["text"],
}
)
seen_contexts.add(context)
title_counts[title] = title_counts.get(title, 0) + 1
if len(passages) == 30:
break
if len(passages) != 30:
raise RuntimeError(
"The Hub sample did not produce 30 passages. "
"Check dataset access and the specified revision."
)
return passages, questions
def make_evaluation_cases(questions):
"""Create answerable and unanswerable evaluation examples."""
if len(questions) != 30:
raise ValueError("Load the Hub sample before creating evaluation cases.")
|