| """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 |
| 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.") |
|
|
