{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Simple Retrieval-Augmented Generation\n", "\n", "Lecture 10 | CMU ANLP Fall 2025 | Instructor: Sean Welleck\n", "\n", "This notebook shows a minimal implementation of Retrieval-Augmented Generation (RAG) for answering questions about the course.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "from transformers import AutoTokenizer, AutoModelForCausalLM\n", "from sentence_transformers import SentenceTransformer\n", "from typing import List, Tuple\n", "import json" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "# Load embedding model\n", "embed_model = SentenceTransformer('Qwen/Qwen3-Embedding-0.6B')\n", "\n", "# Load language model\n", "lm_model_name = \"Qwen/Qwen3-0.6B\"\n", "tokenizer = AutoTokenizer.from_pretrained(lm_model_name)\n", "lm_model = AutoModelForCausalLM.from_pretrained(lm_model_name)\n", "lm_model.eval()\n", "\n", "# Set padding token\n", "tokenizer.pad_token = tokenizer.eos_token" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Document Preparation\n", "\n", "Load pre-parsed documents from JSON file." ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Loaded 43 documents from https://cmu-l3.github.io/anlp-fall2025/\n", "\n", "Example documents:\n", "1. Advanced Natural Language Processing / Fall 2025 Advanced natural language processing is an introductory graduate-level course on natural language processing aimed at students who are interested in doing cutting-edge research in the field. In it, we describe fundamental tasks in natural language processing as well as methods to solve these tasks. The course focuses on modern methods using neural networks, and covers the basic modeling, learning, and inference algorithms required therefore.\n", "2. The class culminates in a project in which students attempt to reimplement and improve upon a research paper in a topic of their choosing.\n", "3. Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya Logistics Class times: TR 2:00pm - 3:20pm Room: TEP 1403 Course identifier: LTI 11-711 Piazza: Piazza Code: GitHub Office hours: Location Day Time Sean Welleck GHC 6513 Tuesday 4:00-5:00 PM Joel Mire WEH 3110 Tuesday 3:30-4:30 PM Chen Wu GHC 5417 Tuesday 4:00-5:00 PM Dareen Alharthi GHC 5417 Monday 10:00-11:00 AM Neel Bhandari GHC 5417 Friday 12:00-1:00 PM Akshita Gupta GHC 5417 Friday 4:00-5:00 PM Ashish Marisetty GHC 5417 Friday 2:00-3:00 PM Manan Sharma GHC 8115 Monday 6:30 PM-7:30 PM Sanidhya Vijayvargiya GHC 5417 Wednesday 12:00-1:00 PM Grading The assignments will be given a grade of A+ (100), A (96), A- (92), B+ (88), B (85), B- (82), or below.\n", "4. The final grades will be determined based on the weighted average of the quizzes, assignments, and project. Cutoffs for final grades will be approximately 97+ A+, 93+ A, 90+ A-, 87+ B+, 83+ B, 80+ B-, etc., although we reserve some flexibility to change these thresholds slightly. Quizzes: Worth 20% of the grade. Your lowest 3 quiz grades will be dropped. Assignments: There will be 4 assignments (the final one being the project), worth respectively 15%, 15%, 20%, 30% of the grade.\n", "5. Course description The course covers key algorithmic foundations and applications of advanced natural language processing. There are no hard pre-requisites for the course, but programming experience in Python and knowledge of probability and linear algebra are expected. It will be helpful if you have used neural networks previously. Acknowledgements.\n", "6. The assignments will be given a grade of A+ (100), A (96), A- (92), B+ (88), B (85), B- (82), or below.The final grades will be determined based on the weighted average of the quizzes, assignments, and project. Cutoffs for final grades will be approximately 97+ A+, 93+ A, 90+ A-, 87+ B+, 83+ B, 80+ B-, etc., although we reserve some flexibility to change these thresholds slightly.Quizzes: Worth 20% of the grade. Your lowest 3 quiz grades will be dropped.Assignments: There will be 4 assignments (the final one being the project), worth respectively 15%, 15%, 20%, 30% of the grade.\n", "7. The final grades will be determined based on the weighted average of the quizzes, assignments, and project. Cutoffs for final grades will be approximately 97+ A+, 93+ A, 90+ A-, 87+ B+, 83+ B, 80+ B-, etc., although we reserve some flexibility to change these thresholds slightly.\n", "8. Quizzes: Worth 20% of the grade. Your lowest 3 quiz grades will be dropped.\n", "9. Assignments: There will be 4 assignments (the final one being the project), worth respectively 15%, 15%, 20%, 30% of the grade.\n", "10. Quiz: There will be a quiz covering the reading material and/or lecture material that you can fill out on Canvas. The quiz will be released by the end of the day of the class (11:59pm) and will be due at the end of the following day (11:59pm).\n" ] } ], "source": [ "with open('documents.json', 'r', encoding='utf-8') as f:\n", " data = json.load(f)\n", "\n", "documents = data['documents']\n", "source_url = data['source_url']\n", "\n", "print(f\"Loaded {len(documents)} documents from {source_url}\")\n", "print(f\"\\nExample documents:\")\n", "for i, doc in enumerate(documents[:5] + documents[-5:]):\n", " print(f\"{i+1}. {doc}\")\n", " " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Simple chunking" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Created 48 chunks\n", "\n", "Example chunk:\n", "Advanced Natural Language Processing / Fall 2025 Advanced natural language processing is an introductory graduate-level course on natural language processing aimed at students who are interested in doing cutting-edge research in the field. In it, we describe fundamental tasks in natural language processing as well as methods to solve these tasks. The course focuses on modern methods using neural networks, and covers the basic modeling, learning, and inference algorithms required therefore.\n", "\n", "Example chunk:\n", "Quiz: There will be a quiz covering the reading material and/or lecture material that you can fill out on Canvas. The quiz will be released by the end of the day of the class (11:59pm) and will be due at the end of the following day (11:59pm).\n" ] } ], "source": [ "def chunk_text(text, chunk_size, overlap):\n", " words = text.split()\n", " chunks = []\n", " for i in range(0, len(words), chunk_size - overlap):\n", " chunk = ' '.join(words[i:i + chunk_size])\n", " if chunk:\n", " chunks.append(chunk)\n", " return chunks\n", "\n", "all_chunks = []\n", "for doc in documents:\n", " chunks = chunk_text(doc, chunk_size=100, overlap=10)\n", " all_chunks.extend(chunks)\n", "\n", "print(f\"Created {len(all_chunks)} chunks\")\n", "print(f\"\\nExample chunk:\\n{all_chunks[0]}\")\n", "print(f\"\\nExample chunk:\\n{all_chunks[-1]}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Embedding and Indexing\n", "\n", "Create embeddings for all chunks." ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Created embeddings with shape: (48, 1024)\n" ] } ], "source": [ "chunk_embeddings = embed_model.encode(all_chunks, prompt_name=\"query\")\n", "print(f\"Created embeddings with shape: {chunk_embeddings.shape}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Retrieval\n", "\n", "Find relevant chunks for a query." ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Query: Who is the instructor?\n", "\n", "Result 1 (score: 0.605):\n", "Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya\n", "\n", "Result 2 (score: 0.488):\n", "Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya Logistics Class times: TR 2:00pm - 3:20pm Room: TEP 1403 Course identifier: LTI 11-711 Piazza: Piazza Code: GitHub Office hours: Location Day Time Sean Welleck GHC 6513 Tuesday 4:00-5:00 PM Joel Mire WEH 3110 Tuesday 3:30-4:30 PM Chen Wu GHC 5417 Tuesday 4:00-5:00 PM Dareen Alharthi GHC 5417 Monday 10:00-11:00 AM Neel Bhandari GHC 5417 Friday 12:00-1:00 PM Akshita Gupta GHC 5417 Friday 4:00-5:00 PM Ashish Marisetty GHC 5417 Friday 2:00-3:00 PM Manan Sharma GHC 8115 Monday\n", "\n", "Result 3 (score: 0.450):\n", "If you don't have much experience with NLP, it will be helpful to consult with the instructors and TAs to learn about how to do the assignments and course project. Because this is a project-based course, we assume that many of the students taking the course will be interested in turning their assignments or project into research papers.\n", "\n" ] } ], "source": [ "def retrieve(query: str, k: int = 3) -> List[Tuple[str, float]]:\n", " query_embedding = embed_model.encode([query], prompt_name=\"query\")[0]\n", " similarities = np.dot(chunk_embeddings, query_embedding)\n", " similarities = similarities / (np.linalg.norm(chunk_embeddings, axis=1) * np.linalg.norm(query_embedding))\n", " top_indices = np.argsort(similarities)[-k:][::-1]\n", " results = [(all_chunks[i], similarities[i]) for i in top_indices]\n", " return results\n", "\n", "# Test retrieval\n", "query = \"Who is the instructor?\"\n", "results = retrieve(query, k=3)\n", "\n", "print(f\"Query: {query}\\n\")\n", "for i, (chunk, score) in enumerate(results):\n", " print(f\"Result {i+1} (score: {score:.3f}):\\n{chunk}\\n\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Generators" ] }, { "cell_type": "code", "execution_count": 41, "metadata": {}, "outputs": [], "source": [ "def generate_with_context(query: str, context_chunks: List[str], max_new_tokens: int = 100) -> str:\n", " context = \"\\n\".join([f\"- {chunk}\" for chunk in context_chunks])\n", " \n", " # Create messages for chat template\n", " messages = [\n", " {\n", " \"role\": \"system\",\n", " \"content\": \"You are a helpful assistant answering questions about a course. Use only the provided context to answer questions. Be concise and accurate. Only generate your answer.\"\n", " },\n", " {\n", " \"role\": \"user\",\n", " \"content\": f\"\"\"Here is some context about the course:\n", "\n", "{context}\n", "\n", "Based on this context, please answer the following question:\n", "{query}\"\"\"\n", " }\n", " ]\n", " \n", " prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)\n", " inputs = tokenizer(prompt, return_tensors=\"pt\", truncation=True, max_length=512)\n", " \n", " # Generate\n", " with torch.no_grad():\n", " outputs = lm_model.generate(\n", " **inputs,\n", " max_new_tokens=max_new_tokens,\n", " temperature=0.1,\n", " top_p=0.95,\n", " do_sample=True,\n", " pad_token_id=tokenizer.eos_token_id\n", " )\n", " \n", " # Decode and extract answer\n", " response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n", " try:\n", " answer = response.split(\"<|im_start|>assistant\")[-1].split(\"\")[-1].strip()\n", " except IndexError:\n", " parts = response.split(query)\n", " answer = parts[-1].strip() if len(parts) > 1 else response\n", " return answer\n", "\n", "def generate_without_context(query: str, max_new_tokens: int = 100) -> str:\n", " \n", " messages = [\n", " {\n", " \"role\": \"system\",\n", " \"content\": \"You are a helpful assistant answering questions about a course, Advanced NLP Fall 2025. Use only the provided context to answer questions. Be concise and accurate. Only generate your answer.\"\n", " },\n", " {\n", " \"role\": \"user\",\n", " \"content\": query\n", " }\n", " ]\n", " \n", " prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)\n", " inputs = tokenizer(prompt, return_tensors=\"pt\")\n", " \n", " # Generate\n", " with torch.no_grad():\n", " outputs = lm_model.generate(\n", " **inputs,\n", " max_new_tokens=max_new_tokens,\n", " temperature=0.1,\n", " top_p=0.95,\n", " do_sample=True,\n", " pad_token_id=tokenizer.eos_token_id\n", " )\n", " \n", " # Decode and extract answer\n", " response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n", " try:\n", " answer = response.split(\"<|im_start|>assistant\")[-1].split(\"\")[-1].strip()\n", " except IndexError:\n", " parts = response.split(query)\n", " answer = parts[-1].strip() if len(parts) > 1 else response\n", " \n", " return answer" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## RAG Pipeline\n", "\n", "Retrieve then generate." ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [], "source": [ "def rag_answer(query: str, k: int = 3) -> str:\n", " results = retrieve(query, k=k)\n", " context_chunks = [chunk for chunk, _ in results]\n", " answer = generate_with_context(query, context_chunks)\n", " return answer, results" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Comparison: With vs Without RAG\n", "\n", "*Do you notice any errors in the RAG outputs? Also try making additional queries, and find ones that lead to errors.*" ] }, { "cell_type": "code", "execution_count": 47, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "============================================================\n", "Query: Who's the instructor?\n", "\n", "WITHOUT RAG:\n", "The instructor is Dr. Emily Carter.\n", "\n", "WITH RAG:\n", "Retrieved chunks:\n", " 1. (score: 0.603) Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya...\n", " 2. (score: 0.487) Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya Logistics Class times: TR 2:00pm - 3:20pm Room: TEP 1403 Course identifier: LTI 11-711 Piazza: Piazza Code: GitHub Office hours: Location Day Time Sean Welleck GHC 6513 Tuesday 4:00-5:00 PM Joel Mire WEH 3110 Tuesday 3:30-4:30 PM Chen Wu GHC 5417 Tuesday 4:00-5:00 PM Dareen Alharthi GHC 5417 Monday 10:00-11:00 AM Neel Bhandari GHC 5417 Friday 12:00-1:00 PM Akshita Gupta GHC 5417 Friday 4:00-5:00 PM Ashish Marisetty GHC 5417 Friday 2:00-3:00 PM Manan Sharma GHC 8115 Monday...\n", "\n", "Generated answer:\n", "Sean Welleck\n", "\n", "============================================================\n", "Query: Who teaches the course?\n", "\n", "WITHOUT RAG:\n", "The course is taught by Dr. Emily Carter.\n", "\n", "WITH RAG:\n", "Retrieved chunks:\n", " 1. (score: 0.648) Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya...\n", " 2. (score: 0.597) Course Details Instructor Sean Welleck Teaching Assistants Joel Mire Chen Wu Dareen Alharthi Neel Bhandari Akshita Gupta Ashish Marisetty Manan Sharma Sanidhya Vijayvargiya Logistics Class times: TR 2:00pm - 3:20pm Room: TEP 1403 Course identifier: LTI 11-711 Piazza: Piazza Code: GitHub Office hours: Location Day Time Sean Welleck GHC 6513 Tuesday 4:00-5:00 PM Joel Mire WEH 3110 Tuesday 3:30-4:30 PM Chen Wu GHC 5417 Tuesday 4:00-5:00 PM Dareen Alharthi GHC 5417 Monday 10:00-11:00 AM Neel Bhandari GHC 5417 Friday 12:00-1:00 PM Akshita Gupta GHC 5417 Friday 4:00-5:00 PM Ashish Marisetty GHC 5417 Friday 2:00-3:00 PM Manan Sharma GHC 8115 Monday...\n", "\n", "Generated answer:\n", "Sean Welleck, Joel Mire, Chen Wu, Dareen Alharthi, Neel Bhandari, Akshita Gupta, Ashish Marisetty, and Manan Sharma.\n", "\n", "============================================================\n", "Query: What is the late policy?\n", "\n", "WITHOUT RAG:\n", "The late policy is that students may submit their work late, but they must complete the assignment by the deadline.\n", "\n", "WITH RAG:\n", "Retrieved chunks:\n", " 1. (score: 0.538) Late Day Policy: In case there are unforeseen circumstances that don't let you turn in your assignment on time, 5 late days total for assignments 1, 2, 3.1, and 3.2 will be allowed. Note that other than these late days, we will not be making exceptions and extending deadlines except for documented health reasons , so please try to be frugal with your late days and use them only if necessary....\n", " 2. (score: 0.385) Assignments that are late beyond the allowed late days will be graded down one third-grade per day late (e.g., A to A- for one day, and A to B+ for two days). Plagiarism/Code Reuse Policy: All assignments are expected to be conducted under the CMU policy for academic integrity . All rules here apply and violations will be subject to penalty including zero credit on the assignment, failing the course, or other disciplinary measures....\n", "\n", "Generated answer:\n", "The late policy allows 5 late days for assignments 1, 2, 3.1, and 3.2, with no exceptions except for documented health reasons.\n", "\n", "============================================================\n", "Query: How much are quizzes worth?\n", "\n", "WITHOUT RAG:\n", "Quizzes are worth 10% of the total grade.\n", "\n", "WITH RAG:\n", "Retrieved chunks:\n", " 1. (score: 0.713) Quizzes: Worth 20% of the grade. Your lowest 3 quiz grades will be dropped....\n", " 2. (score: 0.599) The quiz will be released by the end of the day of the class (11:59pm) and will be due at the end of the following day (11:59pm). Quizzes are worth 20% of the grade. Your lowest 3 quiz grades will be dropped. We provide the drop days in case you have to miss a quiz (e.g., due to travel, unexpected circumstances). Please do not contact the TAs or Instructors about additional quiz drops; we provide three quiz drops to cover unforeseen circumstances....\n", "\n", "Generated answer:\n", "Quizzes are worth 20% of the grade.\n", "\n", "============================================================\n", "Query: When is assignment 3.1 released?\n", "\n", "WITHOUT RAG:\n", "Assignment 3.1 is released on **Friday, April 12, 2025**.\n", "\n", "WITH RAG:\n", "Retrieved chunks:\n", " 1. (score: 0.572) Details of Each Assignment Assignment 1: Build Your Own LLaMa (Individual assignment) Released: Sep 9 Due: Sep 25 Assignment 2: End-to-end NLP System Building (Group assignment) Released: Sep 25 Due: Oct 9 Assignment 3: Project Proposal & State-of-the-art Reimplementation (Group assignment) Assignment 3.1: Literature Review & Project Proposal Released: Oct 9 Due: Oct 30 Assignment 3.2: Baseline Reproduction Released: Oct 9 Due: Nov 13 Assignment 4: Final Project (Group assignment) Released: Oct 9 Due: Dec 9 Details to be provided later.....\n", " 2. (score: 0.570) Details of Each Assignment Assignment 1: Build Your Own LLaMa (Individual assignment) Released: Sep 9 Due: Sep 25 Assignment 2: End-to-end NLP System Building (Group assignment) Released: Sep 25 Due: Oct 9 Assignment 3: Project Proposal & State-of-the-art Reimplementation (Group assignment) Assignment 3.1: Literature Review & Project Proposal Released: Oct 9 Due: Oct 30 Assignment 3.2: Baseline Reproduction Released: Oct 9 Due: Nov 13 Assignment 4: Final Project (Group assignment) Released: Oct 9 Due: Dec 9 Details to be provided later....\n", "\n", "Generated answer:\n", "Oct 9.\n", "\n", "============================================================\n", "Query: What time does the class meet?\n", "\n", "WITHOUT RAG:\n", "The class meets at 9:00 AM.\n", "\n", "WITH RAG:\n", "Retrieved chunks:\n", " 1. (score: 0.594) Logistics Class times: TR 2:00pm - 3:20pm Room: TEP 1403 Course identifier: LTI 11-711 Piazza: Piazza Code: GitHub Office hours: Location Day Time Sean Welleck GHC 6513 Tuesday 4:00-5:00 PM Joel Mire WEH 3110 Tuesday 3:30-4:30 PM Chen Wu GHC 5417 Tuesday 4:00-5:00 PM Dareen Alharthi GHC 5417 Monday 10:00-11:00 AM Neel Bhandari GHC 5417 Friday 12:00-1:00 PM Akshita Gupta GHC 5417 Friday 4:00-5:00 PM Ashish Marisetty GHC 5417 Friday 2:00-3:00 PM Manan Sharma GHC 8115 Monday 6:30 PM-7:30 PM Sanidhya Vijayvargiya GHC 5417 Wednesday 12:00-1:00 PM...\n", " 2. (score: 0.573) on Canvas. The quiz will be released by the end of the day of the class (11:59pm) and will be due at the end of the following day (11:59pm). Questions and Discussion : Ideally in class or through Piazza so we can share information with the class, but coming to office hours is also encouraged....\n", "\n", "Generated answer:\n", "The class meets at TR 2:00pm - 3:20pm.\n", "\n" ] } ], "source": [ "test_queries = [\n", " \"Who's the instructor?\",\n", " \"Who teaches the course?\",\n", " \"What is the late policy?\",\n", " \"How much are quizzes worth?\",\n", " \"When is assignment 3.1 released?\",\n", " \"What time does the class meet?\"\n", "]\n", "\n", "for query in test_queries:\n", " print(\"=\" * 60)\n", " print(f\"Query: {query}\\n\")\n", " \n", " # Without RAG\n", " print(\"WITHOUT RAG:\")\n", " answer_no_rag = generate_without_context(query)\n", " print(f\"{answer_no_rag}\\n\")\n", " \n", " # With RAG\n", " print(\"WITH RAG:\")\n", " answer_rag, retrieved = rag_answer(query, k=2)\n", " \n", " print(\"Retrieved chunks:\")\n", " for i, (chunk, score) in enumerate(retrieved):\n", " print(f\" {i+1}. (score: {score:.3f}) {chunk}...\")\n", " \n", " print(f\"\\nGenerated answer:\")\n", " print(f\"{answer_rag}\\n\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "anlp", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.18" } }, "nbformat": 4, "nbformat_minor": 4 }