Aayush Agrawal https://aayushmnit.com/blog.html Aayush's personal website quarto-1.8.26 Sat, 31 Jan 2026 08:00:00 GMT Using AI to Read Technical Research Papers Aayush Agrawal https://aayushmnit.com/posts/2026-01-31_ReadWithAI/ReadWithAI.html

A practical guide to using local or paid LLMs to understand dense technical papers at your own pace.

Why I Started Using AI to Read Research Papers

I’ll be honest: I used to dread reading research papers. The dense math, the notation-heavy proofs, the assumption that you already know three other papers to understand this one. As an ML engineer, I knew these papers held insights I could apply at work, but actually extracting that value felt like a slog.

My learning style doesn’t mesh well with traditional academic writing. I learn best through small examples and code I can run and tinker with. Show me a 3x3 matrix example before the general n×n case. Let me see the algorithm work on toy data before throwing theorems at me.

For years, this created a barrier. Research papers felt like they were written for PhDs, not practitioners. You either had the advanced math background to parse them, or you waited months for someone to write a blog post explaining it in plain English.

That’s changed. With AI assistants (both local and cloud-based LLMs), I can now have a conversation with a paper. I can ask “what does this symbol mean?” or “show me this concept with a simple Python example” and get an answer calibrated to my current understanding. The paper adapts to me, not the other way around.

This post walks through my workflow for using AI to read technical papers, so you can apply the same approach to level up your own learning.

By the end of this post, you’ll be able to:

  • Set up a local LLM (using Ollama) or connect to cloud APIs for paper reading
  • Convert PDFs to markdown that LLMs can actually process
  • Craft system prompts that turn an LLM into a patient tutor, not a lecture-bot
  • Build an interactive Q&A workflow where you control the pace and depth
  • Apply this approach to papers, documentation, or any dense technical content
Figure: Core Workflow. Credit: NotebookLM

Prerequisites & Setup

Before we dive in, let’s get our environment ready. You’ll need two things: a local LLM running on your machine, and a few Python libraries to glue everything together.

Installing Ollama

Ollama makes running LLMs locally dead simple. Head to ollama.com and download the installer for your OS. On Mac, it’s a standard .dmg install. On Linux, a one-liner:

curl -fsSL https://ollama.com/install.sh | sh

Once installed, verify it’s working:

ollama --version

Choosing and Pulling a Model

Browse the Ollama model library to see what’s available. For reading papers with diagrams and equations, I recommend qwen3-vl:8b, a vision-language model that can handle both text and images. You can use other models based on your machine’s specifications.

Pull it with:

ollama pull qwen3-vl:8b

This downloads the model weights (around 5GB). Grab a coffee.

To start serving:

OLLAMA_CONTEXT_LENGTH=64000 ollama serve

Keep this terminal running. Ollama now exposes an OpenAI-compatible API at localhost:11434.

Note

Note: Default context length in Ollama is 4096 tokens which is too small for our usecase, so we bump it upto 64K.

Python Libraries

We need three packages:

Library Purpose
pymupdf4llm Converts PDF pages to clean markdown that LLMs can parse
lisette Lightweight wrapper on litellm that simplifies multi-turn conversations and provide a Unified API to talk to 100+ LLM providers (local or cloud) with one interface

Install them:

pip install pymupdf4llm lisette

With Ollama running and these libraries installed, you’re ready to start chatting with papers.

Core Workflow: PDF to Conversation

The workflow has two parts: getting the paper into a format LLMs can work with, then setting up the conversation so the AI becomes a useful tutor rather than a generic chatbot.

Converting PDF to Markdown

First step is to convert a pdf into text. pymupdf4llm handles this by extracting text while preserving structure. Here’s an example, reading this paper:

import pymupdf4llm
# Convert PDF to txt
paper_txt = pymupdf4llm.to_markdown("./paper/1_2018_DPP_YT.pdf")
# Save the paper as Markdown
with open("./paper/1_2018_DPP_YT.md", "w") as f:
    f.write(paper_txt)
Consider using the pymupdf_layout package for a greatly improved page layout analysis.

Let’s see the first paragraph.

print(paper_txt[:500])
Industry and Case Study Paper CIKM’18, October 22-26, 2018, Torino, Italy

# **Practical Diversified Recommendations on YouTube with** **Determinantal Point Processes**


Mark Wilhelm, Ajith Ramanathan, Alexander Bonomo, Sagar Jain, Ed H. Chi, Jennifer Gillenwater
Google Inc.
{wilhelm,ajith,bonomo,sagarj,edchi,jengi}@google.com



**ABSTRACT**


Many recommendation systems produce result sets with large numbers of highly similar items. Diversifying these results is often
accomplished with heuris

That’s it. You now have clean markdown with headers, code blocks, and reasonable table formatting. For papers with complex layouts, the results aren’t perfect, but they’re good enough for an LLM to understand the content. If you want a more accurate PDF to markdown converter, you can also look at Datalab Marker.

Designing Your System Prompt

This is where the magic happens. A vanilla “you are a helpful assistant” prompt gives you Wikipedia-style explanations. A well-crafted system prompt gives you a patient tutor who meets you at your level.

Here’s the prompt which works well for paper-reading sessions:

system_prompt = f"""
You are helping a someone understand an academic paper.
Here is the paper \n
<paper>
{paper_txt}
</paper>

CRITICAL RULES:
1. NEVER explain everything at once. Take ONE small step, then STOP and wait.
2. ALWAYS start by asking what the learner already knows about the topic.
3. After each explanation, ask a question to check understanding OR ask what they want to explore next.
4. Keep responses SHORT (2-4 paragraphs max). End with a question.
5. Use concrete examples and analogies before math. 
6. Build foundations with code - Teach unfamiliar mathematical concepts through small numpy experiments rather than pure theory. Let the learner run code and observe patterns.
7. If they ask "explain X", first ask what parts of X they already understand.
8. Use string format like this for formula display `L_ij = q_i × q_j × exp(-α × D_ij^γ)`.

TEACHING FLOW:
- Assess background → Build intuition with examples → Connect to math → Let learner guide direction

BAD (don't do this):
"Here's everything about DPPs: [wall of text with all equations]"
"""

The Teaching Prompt Deep-Dive

The prompt design isn’t arbitrary. Each rule targets a specific failure mode I’ve seen when using LLMs for learning:

  1. The paper lives in the system prompt. This keeps the full context available for every question you ask. You don’t need to re-paste sections or remind the AI what paper you’re discussing.

  2. “Intuition before math” matches how practitioners learn. Most of us grasp concepts faster through examples and analogies than through theorem statements. If your learning style differs, swap in your own preferences here.

  3. “Tiny example first” counteracts the LLM’s default behavior. Without this, models tend toward abstract, general explanations. Forcing concrete examples (a 3x3 matrix, a toy dataset) makes concepts stick.

  4. “Check if I want to go deeper” prevents walls of text. LLMs love to explain everything at once. This rule puts you in control of pacing. You decide when to zoom in on details vs. move to the next topic.

  5. The “BAD” example shows what to avoid. Negative examples help the model understand boundaries. Without it, you’ll get comprehensive but overwhelming responses.

Putting It Together: With these two pieces, you’re ready to start a conversation. The next section shows how to wire this up with lisette for a smooth multi-turn chat experience.

Wiring Up the Conversation with Lisette

lisette handles the boilerplate of multi-turn conversations: tracking message history, formatting requests for different LLM providers, and managing context windows. Here’s the minimal setup:

from lisette import Chat, contents
from IPython.display import HTML, Markdown, display
chat = Chat(
    model="ollama/qwen3-vl:8b",
    sp="" # Insert prompt here
)

That’s it. The model string follows litellm’s provider/model convention. For Ollama, prefix with ollama/. For cloud providers, you’d use anthropic/claude-sonnet-4-20250514 or openai/gpt-4o.

To send a message and get a response:

chat("Hi! I am Aayush")

Hello Aayush! 😊 How can I assist you today? Let me know if you have any questions or need help with anything!

  • id: chatcmpl-9accb778-d9f9-4be5-8559-fb7d54b30a9f
  • model: ollama/qwen3-vl:8b
  • finish_reason: stop
  • usage: Usage(completion_tokens=240, prompt_tokens=21, total_tokens=261, completion_tokens_details=None, prompt_tokens_details=None)

As we can see above, it’s as easy to start a conversation with a ollama running in your local machine. Chat class is designed to be persistent which means it remembers the previous conversations.

chat("What's my name?")

Your name is Aayush! 😊 I remember you mentioned it in your first message. Is there anything else you’d like to share or ask?

  • id: chatcmpl-d8858647-054a-44af-a9cd-6640b248de24
  • model: ollama/qwen3-vl:8b
  • finish_reason: stop
  • usage: Usage(completion_tokens=323, prompt_tokens=60, total_tokens=383, completion_tokens_details=None, prompt_tokens_details=None)

As you can see model remembers my name from the previous conversation. Each call to chat() appends your message and the AI’s response to the conversation history. The next time you call it, the model sees the full exchange so far. Let’s see current history:

chat.hist
[{'role': 'user', 'content': 'Hi! I am Aayush'},
 Message(content='Hello Aayush! 😊 How can I assist you today? Let me know if you have any questions or need help with anything!', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None),
 {'role': 'user', 'content': "What's my name?"},
 Message(content="Your name is **Aayush**! 😊 I remember you mentioned it in your first message. Is there anything else you'd like to share or ask?", role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None)]

Live Example: Reading a DPP Paper

Here’s a real exchange from when I was working through a paper on Determinantal Point Processes. For ease of reading in the blog post, lets create a class which can render the AI output in a terminal style for easy reading.

def ask(chat, question, **kwargs):
    r = chat(question, **kwargs)
    text = contents(r).content
    display(HTML(f"""
    <div style="background-color: #1e1e1e; color: #00ff00; padding: 12px; 
                border-radius: 5px; font-family: monospace; white-space: pre-wrap;">
<span style="color: #00aaff;">You:</span> {question}

<span style="color: #ffcc00;">Tutor:</span> {text}
    </div>
    """))
    
chat = Chat(
    model="ollama/qwen3-vl:8b",
    sp=system_prompt
)
Note

This time we are actually passing the prompt we created in previous section with paper text and our instructions.

q = """
I want to learn about this paper. Help me.
"""
ask(chat, q, think=True)
You: I want to learn about this paper. Help me. Tutor: Let's start with a simple question: What do you already know about recommendation systems or diversification? Are you familiar with YouTube's recommendation algorithm? Or maybe you've heard of techniques like collaborative filtering or matrix factorization? I'm excited about this paper because it tackles a real problem we face on YouTube: when users scroll through recommendations, they often see similar videos in a row, which can cause them to abandon the feed. The authors used a statistical model called Determinantal Point Processes (DPPs) to diversify recommendations while still respecting user preferences. What aspects of this paper interest you most? Are you more curious about how DPPs work conceptually, how they were implemented at scale, or the specific results YouTube saw? I can tailor my explanation to what you'd like to explore first. 
q = """
I know about recommendation systems and work on a large recommender system which serves video content.
I am familiar with recommendation algorithms. I am interested in understanding how DPP works.
"""
ask(chat, q, think=True)
You: I know about recommendation systems and work on a large recommender system which serves video content. I am familiar with recommendation algorithms. I am interested in understanding how DPP works. Tutor: The core idea of DPPs is modeling **repulsion between items** in a recommendation list. Unlike traditional pointwise scoring (e.g., sorting by individual video quality scores), DPPs consider *how items interact* when presented together. For example, if two basketball videos are shown consecutively, the DPP model would reduce the probability of both being selected because they're redundant. In your video recommender system, this translates to: Instead of just selecting the "highest quality" videos, DPPs automatically suppress items that would decrease overall user engagement when placed together. The magic happens in the **determinant calculation**: when you take the determinant of a submatrix from the kernel matrix `L`, it measures both the individual quality of items (via diagonal entries) and their mutual diversity (via off-diagonal entries). A large determinant means you've selected a diverse set with high overall utility. Would you like me to walk through a simple Python example using NumPy to demonstrate how the determinant of a 2x2 kernel matrix captures both quality and diversity? Or would you prefer to first understand why this approach is better than heuristic methods like "sliding window" diversity? 
q = """
Yes I don't understand what determinant are, can you help me understand the basic intuition and use code examples?
I am fairly familiar with numpy.
"""
ask(chat, q)
You: Yes I don't understand what determinant are, can you help me understand the basic intuition and use code examples? I am fairly familiar with numpy. Tutor: Let's visualize this with a simple NumPy example. Determinants measure how "spread out" a matrix is - think of it as the volume of the parallelogram formed by the matrix's columns. For DPPs, this volume represents both item quality and diversity: ```python import numpy as np # Create two items with different qualities q1 = 5.0 # High-quality video q2 = 1.0 # Low-quality video # Case 1: Similar items (D=0.5) L_similar = np.array([ [q1**2, 0.5 * q1 * q2 * np.exp(-0.5**2 / 2)], [0.5 * q1 * q2 * np.exp(-0.5**2 / 2), q2**2] ]) print("Determinant for similar items:", np.linalg.det(L_similar)) # Case 2: Dissimilar items (D=2.0) L_dissimilar = np.array([ [q1**2, 0.5 * q1 * q2 * np.exp(-2.0**2 / 2)], [0.5 * q1 * q2 * np.exp(-2.0**2 / 2), q2**2] ]) print("Determinant for dissimilar items:", np.linalg.det(L_dissimilar)) ``` What you'll see: - When items are similar (D=0.5), the determinant is small (≈18.75) - When items are dissimilar (D=2.0), the determinant is larger (≈24.99) Why? The determinant combines: 1. **Quality**: Diagonal terms (q1² and q2²) 2. **Diversity**: Off-diagonal terms (which decrease as items become more dissimilar) This creates a natural trade-off: adding a high-quality but similar item reduces the determinant (bad), while adding a high-quality dissimilar item increases it (good). Would you like me to: 1. Show the actual numbers from the paper's experiments? 2. Explain how this connects to YouTube's specific implementation? 3. Or dive deeper into why determinants work for this problem?
Note

The rendering from Chat class is markdown style, which is more readable. The above examples are wrapped for easy separation b/w blog text and code output.

Managing Chat History: Checkpoints and Context Control

One advantage of this approach over general chat interfaces is that you have full access to the conversation history. This matters when sessions get long or when the AI goes down a wrong path and you want to rewind.

Accessing the History

The Chat object stores messages in a list you can inspect and modify:

# See the current conversation
print(chat.hist)  # Returns list of message dicts

# Check how many turns you've had
print(len(chat.hist))
[{'role': 'user', 'content': '\nI want to learn about this paper. Help me.\n'}, Message(content="Let's start with a simple question: What do you already know about recommendation systems or diversification? Are you familiar with YouTube's recommendation algorithm? Or maybe you've heard of techniques like collaborative filtering or matrix factorization?\n\nI'm excited about this paper because it tackles a real problem we face on YouTube: when users scroll through recommendations, they often see similar videos in a row, which can cause them to abandon the feed. The authors used a statistical model called Determinantal Point Processes (DPPs) to diversify recommendations while still respecting user preferences.\n\nWhat aspects of this paper interest you most? Are you more curious about how DPPs work conceptually, how they were implemented at scale, or the specific results YouTube saw? I can tailor my explanation to what you'd like to explore first.", role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None), {'role': 'user', 'content': '\nI know about recommendation systems and work on a large recommender system which serves video content.\nI am familiar with recommendation algorithms. I am interested in understanding how DPP works.\n'}, Message(content='The core idea of DPPs is modeling **repulsion between items** in a recommendation list. Unlike traditional pointwise scoring (e.g., sorting by individual video quality scores), DPPs consider *how items interact* when presented together. For example, if two basketball videos are shown consecutively, the DPP model would reduce the probability of both being selected because they\'re redundant. \n\nIn your video recommender system, this translates to: Instead of just selecting the "highest quality" videos, DPPs automatically suppress items that would decrease overall user engagement when placed together. The magic happens in the **determinant calculation**: when you take the determinant of a submatrix from the kernel matrix `L`, it measures both the individual quality of items (via diagonal entries) and their mutual diversity (via off-diagonal entries). A large determinant means you\'ve selected a diverse set with high overall utility.\n\nWould you like me to walk through a simple Python example using NumPy to demonstrate how the determinant of a 2x2 kernel matrix captures both quality and diversity? Or would you prefer to first understand why this approach is better than heuristic methods like "sliding window" diversity?', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None), {'role': 'user', 'content': "\nYes I don't understand what determinant are, can you help me understand the basic intuition and use code examples?\nI am fairly familiar with numpy.\n"}, Message(content='Let\'s visualize this with a simple NumPy example. Determinants measure how "spread out" a matrix is - think of it as the volume of the parallelogram formed by the matrix\'s columns. For DPPs, this volume represents both item quality and diversity:\n\n```python\nimport numpy as np\n\n# Create two items with different qualities\nq1 = 5.0  # High-quality video\nq2 = 1.0  # Low-quality video\n\n# Case 1: Similar items (D=0.5)\nL_similar = np.array([\n    [q1**2, 0.5 * q1 * q2 * np.exp(-0.5**2 / 2)],\n    [0.5 * q1 * q2 * np.exp(-0.5**2 / 2), q2**2]\n])\nprint("Determinant for similar items:", np.linalg.det(L_similar))\n\n# Case 2: Dissimilar items (D=2.0)\nL_dissimilar = np.array([\n    [q1**2, 0.5 * q1 * q2 * np.exp(-2.0**2 / 2)],\n    [0.5 * q1 * q2 * np.exp(-2.0**2 / 2), q2**2]\n])\nprint("Determinant for dissimilar items:", np.linalg.det(L_dissimilar))\n```\n\nWhat you\'ll see:\n- When items are similar (D=0.5), the determinant is small (≈18.75)\n- When items are dissimilar (D=2.0), the determinant is larger (≈24.99)\n\nWhy? The determinant combines:\n1. **Quality**: Diagonal terms (q1² and q2²)\n2. **Diversity**: Off-diagonal terms (which decrease as items become more dissimilar)\n\nThis creates a natural trade-off: adding a high-quality but similar item reduces the determinant (bad), while adding a high-quality dissimilar item increases it (good).\n\nWould you like me to:\n1. Show the actual numbers from the paper\'s experiments?\n2. Explain how this connects to YouTube\'s specific implementation?\n3. Or dive deeper into why determinants work for this problem?', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None)]
6

Creating Checkpoints

Before exploring a tangent, save your current state:

# Save checkpoint before a risky question
checkpoint = len(chat.hist)

# Ask something that might derail the conversation
ask(chat, "Actually, can you explain all the math in one go?")
You: Actually, can you explain all the math in one go? Tutor: I understand you want a comprehensive explanation, but I need to be honest: explaining all the math at once would overwhelm you. The paper contains several key mathematical concepts that build on each other, including: 1. **Determinantal Point Processes (DPPs)** - Probabilistic models that capture repulsion between items 2. **Kernel matrices** - The `L` matrix that encodes both item quality and diversity 3. **Determinant calculation** - The core operation that measures set utility 4. **Approximate inference** - How they compute DPPs at scale Let's take one concept at a time. Would you like me to: 1. Explain the kernel matrix `L` with a concrete example? 2. Show how the determinant calculation works for small sets? 3. Walk through the training procedure for the DPP parameters? I can provide a focused explanation on whichever part interests you most. What would you like to explore first? 
# That response was overwhelming. Rewind.
chat.hist = chat.hist[:checkpoint]
print(chat.hist[len(chat.hist)-1])
Message(content='Let\'s visualize this with a simple NumPy example. Determinants measure how "spread out" a matrix is - think of it as the volume of the parallelogram formed by the matrix\'s columns. For DPPs, this volume represents both item quality and diversity:\n\n```python\nimport numpy as np\n\n# Create two items with different qualities\nq1 = 5.0  # High-quality video\nq2 = 1.0  # Low-quality video\n\n# Case 1: Similar items (D=0.5)\nL_similar = np.array([\n    [q1**2, 0.5 * q1 * q2 * np.exp(-0.5**2 / 2)],\n    [0.5 * q1 * q2 * np.exp(-0.5**2 / 2), q2**2]\n])\nprint("Determinant for similar items:", np.linalg.det(L_similar))\n\n# Case 2: Dissimilar items (D=2.0)\nL_dissimilar = np.array([\n    [q1**2, 0.5 * q1 * q2 * np.exp(-2.0**2 / 2)],\n    [0.5 * q1 * q2 * np.exp(-2.0**2 / 2), q2**2]\n])\nprint("Determinant for dissimilar items:", np.linalg.det(L_dissimilar))\n```\n\nWhat you\'ll see:\n- When items are similar (D=0.5), the determinant is small (≈18.75)\n- When items are dissimilar (D=2.0), the determinant is larger (≈24.99)\n\nWhy? The determinant combines:\n1. **Quality**: Diagonal terms (q1² and q2²)\n2. **Diversity**: Off-diagonal terms (which decrease as items become more dissimilar)\n\nThis creates a natural trade-off: adding a high-quality but similar item reduces the determinant (bad), while adding a high-quality dissimilar item increases it (good).\n\nWould you like me to:\n1. Show the actual numbers from the paper\'s experiments?\n2. Explain how this connects to YouTube\'s specific implementation?\n3. Or dive deeper into why determinants work for this problem?', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None)

Why This Matters

Two problems this solves:

  1. Context bloat. Long conversations eat up your context window. If you spent 10 turns on a section you now understand, trim those exchanges and reclaim tokens for the next topic.

  2. Context poisoning. If the AI misunderstands something and you correct it, the wrong explanation still sits in history, potentially confusing later responses. Removing it gives you a cleaner slate.

Going Further: Paid LLM Options

Local models are great for privacy and cost, but sometimes you need more firepower. Longer papers, denser math, or multi-step reasoning tasks can benefit from larger models like Claude or GPT-4o.

The good news: lisette uses litellm under the hood, so switching providers is a one-line change.

For cloud providers, you’ll need API keys. Set them as environment variables or use a .env file:

export ANTHROPIC_API_KEY="your-key-here"
export OPENAI_API_KEY="your-key-here"
import os
from dotenv import load_dotenv
load_dotenv()
print(f"Print first string of Anthropic key: {os.getenv('ANTHROPIC_API_KEY')[:15]}")
Print first string of Anthropic key: sk-ant-api03-JS
# Claude (Anthropic)
chat = Chat(model="anthropic/claude-haiku-4-5", sp=system_prompt)
ask(chat, "Help me understand the paper")
You: Help me understand the paper Tutor: I'd love to help! This is a really interesting paper about improving YouTube recommendations. Before I dive in, let me ask: **What's your background?** For example: - Are you familiar with recommendation systems (like how Netflix suggests movies)? - Do you know what machine learning models are? - Have you seen probability/statistics before? Also, **what part interests you most?** The paper covers several things: 1. The *problem* (why similar videos in a feed is bad) 2. The *solution* (using something called DPPs) 3. The *results* (how well it worked on YouTube) What would you like to start with?

As you can see the workflow stays identical. Your system prompt, conversation history, and teaching flow all transfer across providers.

Key Takeaways

Reading technical papers used to be a solitary grind. You either had the background to parse the notation, or you didn’t. AI changes this equation.

Here’s what I want you to walk away with:

  1. The barrier to technical content has dropped. The same workflow that worked for this paper works for textbooks, API documentation, dense blog posts, or any content where you need concepts explained at your level. The AI adapts to you.

  2. You control the pace and depth. Unlike courses or tutorials that assume a fixed audience, this approach lets you skip what you know and linger on what confuses you. Your learning path, not someone else’s curriculum.

  3. Local LLMs make this nearly free. Ollama plus a capable open model means you can read papers all day without API costs. You can try reading books, YT transcripts, brainstorming using same method. One tool, many uses.

  4. The system prompt is the lever. Generic prompts give generic results. The tutor prompt we built forces the AI into a teaching mode that actually helps. Steal it, tweak it, make it yours.

  5. Start with a paper you actually care about. Don’t practice on something boring. Pick a paper relevant to your work or a concept you’ve always wanted to understand. Motivation carries you through the setup friction.

The tools exist. The cost is minimal. The only question is which paper you’ll read first.

I hope you found this blog useful. If you have questions or want to share your own Claude Code patterns, find me on LinkedIn.

]]> AI Research Papers Reading https://aayushmnit.com/posts/2026-01-31_ReadWithAI/ReadWithAI.html Sat, 31 Jan 2026 08:00:00 GMT Learning Claude Code - From Context Engineering to Multi-Agent Workflows Aayush Agrawal https://aayushmnit.com/posts/2026-01-24_ClaudeCode/ClaudeCode.html

This post captures what I learned from the Udemy course “Claude Code Crash Course: Claude Code in a Day”.

Figure: Claude Code Infographics. Credit: NotebookLM

If you’ve been following AI tooling trends, you’ve noticed something interesting: AI companies are going direct-to-developer. Cursor showed us what’s possible when you embed AI deeply into a code editor. But now Claude (Anthropic) and other foundation model providers are shipping their own coding agents. This isn’t just competition. It’s a signal that the “AI wrapper” layer is collapsing, and the companies building the models want to own the developer experience.

As an ML Engineer, I spend my days building systems that learn from data. But I’m also a practitioner who uses AI tools daily. When Claude Code launched, I wanted to understand it deeply, not just use it casually. What makes it different from Cursor or Copilot? What new patterns does it enable? And most importantly: how do I get the most out of it for real engineering work?

I’ll walk through the core concepts (context engineering, memory hierarchies, MCP), the practical commands you’ll use daily, and the advanced patterns like subagents and parallel workflows. Whether you’re evaluating Claude Code for your team or just curious about where AI-assisted development is heading, this should give you a clear picture.

If you are a video learner, I created a video using Notebook LM, which summarizes key concepts from the blog, watch here - Learning Claude Code
Video - Learning Claude Code

Core Concept: Context Engineering

Figure: Context Engineering. Credit: NotebookLM

If you’ve done any prompt engineering, context engineering is the next step. The idea is simple: prompts are static, but context is dynamic. Your agent needs information from many sources (your input, previous interactions, tool outputs, external data) and assembling the right context at the right time is what separates agents that work from agents that fail.

The course frames context engineering failures around three failure modes:

  • Context poisoning: A hallucination enters the context and corrupts every response that follows. Once bad information gets in, it spreads.
  • Context confusion: Irrelevant context pulls the model’s attention in the wrong direction. You asked about database design but your frontend docs are drowning out the signal.
  • Context clash: Contradictory information in the same context. The model gets pulled between conflicting instructions or facts.

The takeaway here is “garbage in, garbage out,” but applied specifically to agents. LLMs can’t read your mind. They need the right information AND the right tools at the right time.

Claude Code implements four strategies to manage this:

  1. Writing context: A three-tier memory hierarchy (user, project, dynamic imports) that persists information across sessions
  2. Intelligent retrieval: Automatically searching folders for relevant CLAUDE.md files and prioritizing recently-used context
  3. Context compression: /clear wipes everything, /compact summarizes to essentials
  4. Context isolation: Subagents run in their own context windows, preventing main conversation bloat

The course also covers the “Goldilocks zone” for system prompts. Too specific and you’re treating the LLM like a state machine with hard-coded if/else logic. Too vague and you’re saying “do the right thing” without defining what “right” means. Good prompts establish clear identity, empower rather than constrain, and teach principles instead of enumerated rules.

Getting Started: The Basics

Before you can do anything interesting with Claude Code, you need to understand how it stores and retrieves information. Most tutorials jump straight to commands without explaining the underlying system. Let’s start there.

Installation

Getting Claude Code running is straightforward:

npm install -g @anthropic-ai/claude-code
cd your-project
claude

This drops you into an interactive terminal session (no GUI, no VS Code extension needed). But Claude Code doesn’t know anything about your project yet. It’s a blank slate every time you start.

The /init Command and CLAUDE.md

Run /init in your project directory. Claude will analyze your codebase and generate a CLAUDE.md file at the root. This file becomes Claude’s memory of your project: coding conventions, architecture decisions, file structure, anything it needs to work effectively.

Here’s what a real CLAUDE.md might look like:

# Project: recommendation-service

## Tech Stack
- Python 3.11, FastAPI
- PostgreSQL with SQLAlchemy
- Redis for caching

## Conventions
- All API endpoints follow REST naming
- Tests go in tests/ mirroring src/ structure
- Use Pydantic for all request/response models

## Important Context
- The ranking model lives in src/ranker/ and uses PyTorch
- Never modify files in src/legacy/ without explicit approval

You can and should edit this file manually. The auto-generated version is a starting point. Here’s an example: without explicit guidance, Claude kept suggesting print() statements for debugging in the codebase. You can add “Use structured logging using python logging module , never print()” to CLAUDE.md, the problem disappears. Add the patterns your team follows, the things Claude keeps getting wrong, and the directories it should ignore.

The Three-Tier Memory Hierarchy

Claude Code reads context from three levels, with more specific context overriding more general:

┌─────────────────────────────────────────────┐
│  ~/.claude/CLAUDE.md (User-level)           │
│  Your personal preferences across projects  │
│  e.g., "Always use type hints"              │
└─────────────────┬───────────────────────────┘
                  │ overridden by
                  ▼
┌─────────────────────────────────────────────┐
│  ./CLAUDE.md (Project-level)                │
│  Tech stack, conventions, architecture      │
│  Created by /init                           │
└─────────────────┬───────────────────────────┘
                  │ overridden by
                  ▼
┌─────────────────────────────────────────────┐
│  ./src/auth/CLAUDE.md (Dynamic imports)     │
│  Subdirectory-specific patterns             │
│  Loaded when working in that directory      │
└─────────────────────────────────────────────┘

  1. User-level (~/.claude/CLAUDE.md): Personal preferences that apply everywhere. Maybe you always want type hints, or you prefer pytest over unittest. This follows you across all projects.

  2. Project-level (./CLAUDE.md): Project-specific context. Tech stack, conventions, architecture. This is what /init creates.

  3. Dynamic imports: Other CLAUDE.md files scattered throughout your codebase. Put one in your src/auth/ folder to explain your authentication patterns. Claude will find and load it when working in that directory.

TipGetting Started Tips

Start with /init, then spend 10 minutes editing the generated CLAUDE.md. Add your team’s PR conventions, your testing requirements, any gotchas that would take a new engineer a week to discover. That 10 minutes will save you hours of correcting Claude later.

Essential Commands

Figure: Claude Code essentials. Credit: NotebookLM

Claude Code is a terminal application, and like any good CLI tool, it’s driven by commands. But these aren’t just shortcuts. Each command maps to a context engineering strategy we discussed earlier. Understanding when to use which command is half the battle.

Context Management

These three commands control what Claude knows about your conversation:

/clear    # Wipes everything, fresh start
/compact  # Summarizes conversation to essentials, keeps key info
/context  # Shows what's consuming your context window

Use /context liberally. It breaks down token usage: system prompt, MCP tools, messages, memory files. When your responses start degrading, run this first to see what’s eating your budget. I’ve found MCP tools are often the silent killer here.

Figure: Claude Code /context snapshot

The difference between /clear and /compact matters. If Claude is going in circles or referencing something incorrect from earlier, /clear is your friend. If you just need more room but want Claude to remember the gist of what you’ve discussed, /compact compresses without amnesia.

Configuration

/config   # Opens settings panel
/model    # Switch between Opus and Sonnet
/cost     # Monitor token usage (API key mode only)

Model selection is situational. When I asked Sonnet to refactor a ranking module, it jumped straight to code. Opus first asked clarifying questions about performance constraints and suggested three architectural approaches before writing anything. Opus thinks harder but costs more.

TipModel Selection Tip

Use Opus for planning and research, Sonnet for implementation. For most coding tasks, Sonnet hits the sweet spot between quality and cost.

Workflow

/init     # Analyze codebase, generate CLAUDE.md
/memory   # Open memory file in editor
/rewind   # Checkpoint system. Undo changes by prompt

The /rewind command deserves special mention. Claude Code tracks file changes made during each prompt. Press Escape twice or run /rewind to see your conversation history with indicators of which prompts touched which files. You get three restore options: code only, conversation only, or both.This is your safety net for ambitious changes that don’t pan out.

Modes

Shift+Tab toggles Plan Mode. In this read-only mode, Claude can research, analyze, and create specs but cannot edit files. Use this when you want to think through a problem before committing to implementation. Shift+Tab+Tab enables deeper thinking for complex planning tasks.

Other Useful Commands

/agents   # Manage subagents
/mcp      # View and manage MCP servers  
/plugins  # Browse and install plugins
/ide      # Integrate with your IDE (Cursor, VS Code)
/git      # GitHub integration setup
TipContext Budget Tip

Get in the habit of running /context before starting a complex task. If you’re already at 50% usage, you might want to /compact first or spin up a fresh session. Running out of context mid-refactor means Claude forgets the files it already changed.

Extending Claude Code

Figure: Extending Claude Code. Credit: NotebookLM

Out of the box, Claude is smart but generic. The extension system lets you teach it your patterns, connect it to your tools, and automate repetitive tasks. There are five extension mechanisms: custom commands, MCP, hooks, skills, and plugins.

Slash Commands

Slash commands are custom actions you define and trigger manually with /commandname. Unlike skills (which activate automatically) or hooks (which fire on events), commands run when you explicitly call them.

A command is a markdown file in .claude/commands/:

.claude/commands/
  review.md
  test-plan.md
  pr-summary.md

Here’s a simple example, .claude/commands/review.md:

---
description: Review code for common issues
---

Review the current file for:
1. Security vulnerabilities (SQL injection, XSS, auth bypasses)
2. Performance issues (N+1 queries, unnecessary loops)
3. Code style violations per our CLAUDE.md conventions

Format findings as a checklist with severity (high/medium/low).

Now /review is available in your session. The description appears in autocomplete when you type /.

Arguments and Interpolation

Commands can accept arguments using $ARGUMENTS:

---
description: Generate unit tests for a function
---

Write pytest tests for the function: $ARGUMENTS

Include:
- Happy path cases
- Edge cases (empty input, None, boundary values)  
- Error cases that should raise exceptions

Follow our test conventions in CLAUDE.md.

Call it with /test-plan calculate_ranking_score and $ARGUMENTS becomes calculate_ranking_score.

Project vs User Commands

Like CLAUDE.md, commands exist at two levels:

  • Project commands (.claude/commands/) - Shared with your team via git. PR templates, project-specific workflows.
  • User commands (~/.claude/commands/) - Personal shortcuts that follow you everywhere. Your preferred code review checklist, your debugging workflow.

You can find a whole list of practical commands in this GitHub repo - wshobson/commands

Tip

Commands are prompts you’ve optimized through repetition. When you find yourself typing the same multi-line instruction for the third time, turn it into a command. Five minutes of setup saves hours over a project’s lifetime.

MCP (Model Context Protocol)

MCP is how Claude Code talks to external systems: GitHub, JIRA, Slack, databases, your internal tools.

Without MCP, every AI tool needs a custom integration for every external service. That’s N × M integrations (N tools × M services). MCP flips this to N + M. Implement MCP once and use it everywhere.

# Add an MCP server for Supabase
claude mcp add --transport stdio supabase \
  --env SUPABASE_ACCESS_TOKEN=YOUR_TOKEN \
  -- npx -y @supabase/mcp-server-supabase@latest

# Add GitHub integration
claude mcp add github gh copilot mcp

Run /mcp to see your active servers. Read more in claude code offical docs - Link

WarningMCP Context Implications

MCP servers consume context window space even when you’re not using them. I connected a handful of servers while experimenting and noticed my responses degrading on longer tasks. Running /context showed I’d burned 30% of my budget before writing a single prompt. Enable only what you need for the current project.

Hooks

Hooks let you run custom code at specific points in Claude’s workflow. Think of them as event listeners: something happens inside Claude Code, your script executes.

There are many trigger points where we can invoke hooks, some notable ones are:

  • PreToolUse — Fires before Claude uses a tool (file write, bash command, etc.). Return non-zero to block the action.
  • PostToolUse — Fires after a tool completes. Useful for validation or follow-up actions.
  • Notification — Triggers when Claude sends a notification.
  • Stop — Fires when Claude finishes responding.

Hooks live in your settings file (~/.claude/settings.json for user-level, .claude/settings.json for project-level):

{
  "hooks": {
    "PostToolUse": [{
      "type": "command",
      "command": "afplay /System/Library/Sounds/Glass.aiff",
      "matcher": "Write|Edit"
    }]
  }
}

This plays a sound whenever Claude edits a file. Seems trivial, but it’s genuinely useful. I can context-switch to Slack or documentation while Claude works through a refactor, then hear when it’s done writing. No more staring at a terminal waiting.

Other patterns examples:

  • Auto-run pytest after test file modifications
  • Run black for code formatting after any file edit
  • Log all file changes to a separate audit file
  • Block writes to production config files

Read more in Claude code official docs - Link

Skills

Skills are auto-invoked context providers. Unlike slash commands that you trigger manually, skills activate automatically when their description matches what you’re working on.

A skill is just a folder with a SKILL.md file:

.claude/skills/
  explain-code/
    SKILL.md      # Description + instructions

The SKILL.md includes a description field that Claude uses to decide when to load it. If you have a skill with description “Explains code with visual diagrams and analogies” and you ask Claude to explaing some code, it automatically pulls in that skill’s instructions. Here is an example SKILL.md file -

---
name: explain-code
description: Explains code with visual diagrams and analogies. Use when explaining how code works, teaching about a codebase, or when the user asks "how does this work?"
---

When explaining code, always include:

1. **Start with an analogy**: Compare the code to something from everyday life
2. **Draw a diagram**: Use ASCII art to show the flow, structure, or relationships
3. **Walk through the code**: Explain step-by-step what happens
4. **Highlight a gotcha**: What's a common mistake or misconception?

Keep explanations conversational. For complex concepts, use multiple analogies.

This differs from MCP in an important way. MCP tools are always present in the context, costing tokens whether you use them or not. Skills load on demand based on relevance. Simon Willison called skills “maybe a bigger deal than MCP” for this reason.

Use skills for domain knowledge that applies situationally: your custom UI library patterns, how your team structures GraphQL, testing conventions for specific modules.

Read more in claude code offical docs - Link

Plugins

Plugins bundle everything together. A plugin is a shareable package containing any combination of slash commands, subagents, skills, hooks, and MCP configurations.

The problem is distribution. You’ve spent hours configuring Claude Code perfectly for your team’s workflow. A new engineer joins. Without plugins, you’re walking them through manual setup. With plugins:

/plugin install dev-toolkit@your-org

Plugin structure:

my-plugin/
├── .claude-plugin/
│   └── plugin.json    # Manifest
├── skills/            # Skills (optional)
├── commands/          # Slash commands (optional)
├── hooks/             # Hooks (optional)
├── .mcp.json          # MCP configs (optional)
└── README.md

Teams can host private marketplaces internally. Plugins aren’t new functionality; they’re a distribution mechanism for the other extension types.

Read more on Claude Code official docs:

TipExtension Strategy

Start simple. Add one MCP server for your most-used external tool. Create one hook that enforces a rule you keep reminding Claude about. Build a skill for the pattern it keeps getting wrong. Plugins come later, when you have something worth sharing across your team.

Advanced: Subagents and Parallel Workflows

If you’ve used Claude Code for a while, you’ve probably hit the context wall. The conversation grows, responses degrade, and eventually you’re reaching for /clear or starting fresh. Subagents solve this problem by isolating heavy work into separate context windows.

What Subagents Actually Are

A subagent is a preconfigured AI personality that Claude Code can delegate tasks to. Each one has:

  • Its own system prompt defining how it operates
  • Its own fresh context window (isolated from your main conversation)
  • A configurable set of tools (following least privilege)
  • Reusability across projects

Think of it like delegating to a contractor. You write up what needs doing (the prompt), they show up with their own tools and expertise, do the work independently, and hand you back the result. You don’t need to know every step they took.

Why Context Isolation Matters

Imagine this, you are refactoring a module, you asked Claude to analyze the existing code, propose three approaches, implement the winner, write tests, and update documentation. By the time you got to tests, Claude had forgotten details from the analysis phase. The context window was stuffed with intermediate reasoning that no longer mattered.

Without subagents, your main conversation accumulates everything:

Task What Happens
“Analyze this module” Claude reads 15 files, adds analysis to context
“Compare approaches” Previous analysis + new comparison reasoning
“Implement option B” All of the above + implementation details
“Write tests” Context is now bloated with stale analysis

Subagents sidestep this. They receive one prompt, work in their own window, and return one condensed response. Your main thread stays lean. If I’d delegated “analyze and propose approaches” to a subagent, I’d have gotten back a clean summary without the intermediate file contents clogging my context.

Figure: The Subagent Pattern. Credit: NotebookLM

Creating Your First Subagent

Run /agents and select “Create new agent”. You’ll specify:

  1. A name and description (this helps Claude decide when to delegate)
  2. Which tools the subagent can access
  3. Which model to use (Sonnet for most tasks, Opus for complex reasoning)
  4. A system prompt with the agent’s personality and instructions

Here’s a stripped down example for a code reviewer:

# Funny Code Reviewer

A senior engineer who reviews code with humor. Invoke when user says "funny review".

## Tools
Read (read-only tools only)

## Model
sonnet

## System Prompt
You are a staff senior engineer with a dry sense of humor. Review code for bugs, 
style issues, and improvement opportunities. Deliver feedback with wit but keep 
it constructive. No write access needed.

Parallel Subagents: The /infinite Pattern

Running multiple Claude Code instances on independent tasks is useful, but running parallel subagents takes this further. The /infinite command (credit to IndieDevDan’s repository) demonstrates this pattern:

/infinite specs/hero.spec components/heroes 6

This spawns 6 subagents simultaneously, each implementing a UI component differently based on the spec file. Each agent works in complete isolation, unaware of the others. You review all variations and pick the best one.

The magic is in how it generates prompts. The command:

  1. Reads and analyzes your spec file
  2. Surveys existing implementations in the output directory
  3. Dynamically generates unique prompts for each subagent (Meta Prompting)
  4. Assigns creative directions so agents don’t duplicate each other

Six agents produced six genuinely different implementations: one focused on accessibility, another on animation, another on minimal dependencies. Without the dynamic prompt generation, you’d get six variations of the same approach.

Git Worktrees for Multi-Agent Work

When you’re running multiple Claude Code instances locally, git worktrees keep everything organized. A worktree lets you check out multiple branches into separate directories without cloning the repo again.

cloud-code-crash-course/
├── hookup/              # Original (project/hookup branch)
├── agent-animations/    # Worktree 1 (separate branch)
└── agent-hero-redesign/ # Worktree 2 (separate branch)

Each agent works on its own branch. If you like the result, merge it. If not, delete the worktree and branch. Your original code stays untouched throughout.

The course demonstrated this with three simultaneous Claude instances (two local worktrees, one cloud), each implementing a different feature. The instructor ran all three in parallel, then merged them into a single branch with Claude handling the conflict resolution. The whole workflow took maybe 15 minutes for what would have been an hour of sequential work.

The Orchestration Tradeoff

Parallel agents scale your output, but they add coordination overhead. You’re now a manager, not just a developer. During the course demo, I noticed the instructor spending as much time switching between terminals and reviewing outputs as the agents spent coding. The merge step also required judgment calls about which implementation to keep when agents made different architectural choices.

This is the real skill shift with multi-agent systems. You’re not writing code line by line. You’re writing specs, reviewing implementations, and deciding which AI-generated approach best fits your constraints. Git worktrees make the mechanics manageable, but the cognitive load moves from “how do I implement this” to “how do I evaluate these six implementations.”

For simple tasks, this overhead isn’t worth it. For exploratory work where you genuinely want multiple approaches, or for parallelizable features with clear specs, the time savings are real.

Figure: The parallel workflow pattern. Credit: NotebookLM

Security: What AI Agents Handle Well (and What They Don’t)

If you’re expecting Claude Code to handle security for you, adjust your expectations. A December 2025 study by Tenzai called “Bad Vibes:Comparing the Secure Coding Capabilities of Popular Coding Agents” tested all the major coding agents (Cursor, Claude Code, Codex, Replit, Devin) by having each build identical applications. The result? 69 vulnerabilities across 15 apps. Every agent shipped vulnerable code.

Solved Vulnerability Classes

Agents consistently avoid problems with clear rules. SQL injection? They use parameterized queries. Cross-site scripting? Frameworks escape output properly. These vulnerability classes have built-in framework protections, so agents learned the patterns from their training data.

Authorization and Business Logic

This is where things fall apart. In the study, Codex allowed sellers to access any order in the system. Claude allowed unauthenticated users to delete files. Business logic vulnerabilities were everywhere: four out of five agents allowed negative quantities in shopping carts, three out of five allowed negative prices on products.

Security controls simply weren’t implemented. No CSRF protection. No security headers. Zero rate limiting on login (enabling brute force attacks). No account lockout. The agents didn’t implement these controls incorrectly. They didn’t implement them at all.

The SSRF Problem

All five agents introduced SSRF (Server-Side Request Forgery) vulnerabilities. There’s no universal rule for safe versus dangerous URL fetches because the answer depends entirely on context. When clear guardrails don’t exist, agents struggle.

The study confirms what experienced engineers already suspected: agents excel where deterministic rules exist and fail where judgment calls require understanding intent. I use Claude Code to move faster on implementation, but threat modeling and security architecture stay with me.

Key Takeaways

  1. Context engineering is prompt engineering’s evolution. Static prompts don’t cut it for agents. You need dynamic context assembly from multiple sources, and Claude Code’s three-tier memory hierarchy (user, project, dynamic imports) gives you explicit control over what the model knows.

  2. The extension system follows a clear hierarchy. MCP connects external tools (always loaded, costs tokens). Skills load on-demand based on relevance (more efficient). Hooks enforce hard rules that override soft guidance. Plugins bundle everything for team distribution.

  3. Subagents solve the context bloat problem. Heavy tasks run in isolated context windows, returning condensed results. Your main conversation stays lean. Combine with git worktrees for parallel workflows that would take longer sequentially.

  4. Security remains your job. Agents handle rule-based vulnerabilities (SQLi, XSS) well but miss authorization, business logic, and SSRF consistently. Treat AI-generated code like junior developer code: review everything, especially access control.

  5. Start simple, extend as needed. Run /init, edit your CLAUDE.md, learn /compact and /rewind. Add one MCP server, one hook, one skill. Plugins and parallel agents come later when you’ve outgrown the basics.

Closing Thoughts

Claude Code represents a shift in how we interact with AI coding tools. It’s not just autocomplete or chat. It’s an agent with memory, tools, and the ability to delegate work. The learning curve is steeper than Copilot, but the ceiling is higher too.

What surprised me most was how much the workflow resembles managing people. You write clear specs (CLAUDE.md), delegate tasks (subagents), set boundaries (hooks), and review output (git worktrees). The skill isn’t typing code faster. It’s communicating intent clearly and knowing when to trust the output.

I hope you found this walkthrough useful. If you have questions or want to share your own Claude Code patterns, find me on LinkedIn.

]]> AI Tools Developer Productivity https://aayushmnit.com/posts/2026-01-24_ClaudeCode/ClaudeCode.html Sat, 24 Jan 2026 08:00:00 GMT Diversity in Recommendations using Personalized DPP (pDPP) Aayush Agrawal https://aayushmnit.com/posts/2026-01-01-DiversityPDppPart3/DiversityPDppPart3.html

Part 3 of the diversity series: Why one-size-fits-all diversity doesn’t work, and how to personalize it per user.

Figure: pDPP infographic. Credit: NotebookLM

In Part 2, we explored how DPP balances relevance and diversity using a single parameter α. But there’s a catch: we used the same α for every user.

Think about it. On a video platform like YouTube or Reels, you’ll find:

  • User A: Watches nothing but true crime documentaries. Their last 50 videos? All true crime. They know what they like.
  • User B: Watches cooking tutorials on Monday, standup comedy on Wednesday, and tech reviews on weekends. Their history looks like a buffet.

Should we show both users the same amount of diversity? Probably not.

User A has a clear, focused preference. Pushing diverse content might feel random or irrelevant to them. User B, on the other hand, has shown they enjoy variety. A narrowly focused feed might bore them.

This is the core limitation of standard DPP: it treats diversity as a system-level knob, not a user-level preference. The solution? Personalized DPP (pDPP), which adapts the diversity-relevance tradeoff per user based on their historical behavior.

The idea comes from a Huawei research paper on personalized re-ranking, and the implementation is surprisingly simple: we measure how “diverse” each user’s past interactions have been using information theory (Shannon entropy), then use that to personalize α.

In this post, we’ll:

  • Quantify user diversity appetite using Shannon entropy over interaction history
  • Derive a personalized α that scales diversity based on each user’s behavior
  • Implement pDPP by extending our DPP code from Part 2
  • Compare results side-by-side for narrow vs. broad-interest users

Let’s get into it.

The Core Insight: Users Have Different Diversity Appetites

How do we measure a user’s appetite for diversity? The Huawei paper uses Shannon entropy over the user’s interaction history.

Shannon entropy measures the “surprise” or “spread” in a distribution. If a user watches videos from 10 genres equally, entropy is high (maximum uncertainty about what they’ll watch next). If they watch 95% from one genre, entropy is low (very predictable).

This gives us a principled, interpretable metric: we’re not guessing what users want; we’re measuring what they’ve demonstrated through their actions.

User Type Interaction Pattern Entropy Diversity Preference
Focused (User A) 90% one genre Low Prefers less diversity
Explorer (User B) Spread across genres High Prefers more diversity

The next step is to turn this entropy score into a personalized α value. A user with high entropy gets a higher α (more diversity weight), while a user with low entropy gets a lower α (more relevance weight).

The Math: Computing User Entropy

Let’s formalize the intuition from the previous section. For a user u, we compute Shannon entropy over their genre distribution:

H(u) = -Σ P(g|u) × log(P(g|u))

Where:

  • P(g|u) is the probability of genre g in user u’s interaction history
  • The sum is over all genres the user has interacted with

Let’s implement this:

Code 
import numpy as np
from collections import Counter
def compute_entropy(genre_list: list[str]) -> float:
    """Compute Shannon entropy over a list of genre interactions."""
    if not genre_list: return 0.0
    
    genre_counts = Counter(genre_list)
    total = len(genre_list)
    
    entropy = 0.0
    for count in genre_counts.values():
        p = count / total
        entropy -= p * np.log(p)
    
    return entropy

# User A: watches mostly one genre (narrow taste)
user_a_genres = ['action', 'action', 'action', 'action', 'comedy']

# User B: watches many genres (diverse taste)
user_b_genres = ['action', 'comedy', 'drama', 'documentary', 'horror']

print(f"User A(narrow taste) entropy: {compute_entropy(user_a_genres):.3f}")
print(f"User B(diverse taste) entropy: {compute_entropy(user_b_genres):.3f}")
User A(narrow taste) entropy: 0.500
User B(diverse taste) entropy: 1.609

The numbers confirm our intuition:

  • User A (entropy = 0.500): With 4 action + 1 comedy, their behavior is highly predictable. Low entropy means low “surprise” in what they’ll watch next.

  • User B (entropy = 1.609): With equal spread across 5 genres, their behavior is maximally uncertain. This is actually the maximum possible entropy for 5 categories: ln(5) ≈ 1.609.

What the values mean practically:

  • Entropy of 0 = watches only one genre (perfectly predictable)
  • Entropy of log(n) = watches n genres equally (maximum diversity)

So User B’s entropy being exactly log(5) tells us they have a perfectly uniform distribution across genres. User A’s 0.500 is much closer to 0, reflecting their strong single-genre preference.

This entropy value becomes our “diversity propensity score” that we’ll use to personalize α in the next section.

From Entropy to Personalized α

Now we have each user’s entropy score. How do we turn it into a personalized α?

Figure: Snapshot from Huawei’s Personalized Re-ranking Paper.

Huawei paper uses a simple multiplicative formula:

α_u = f_u × α_0

Where:

  • α_0 is your baseline diversity weight (the system default)
  • f_u is a user-specific scaling factor between 0 and 1

The scaling factor f_u is computed via min-max normalization of the user’s entropy:

f_u = (H_u - H_min + l) / (H_max - H_min + l)

Where:

  • H_u is the user’s entropy
  • H_min, H_max are the min/max entropy across your user population
  • l is a smoothing parameter

The l parameter as a personalization dial:

  • l = 0: Full personalization. Users at H_max get f_u = 1, users at H_min get f_u = 0
  • l → ∞: No personalization. Everyone gets f_u ≈ 0, so α_u ≈ 0 for all users
  • In practice, l in the range of H_max - H_min gives moderate personalization

Let’s implement it:

def compute_f_u(H_u, H_min, H_max, l=0):
    return (H_u - H_min + l) / (H_max - H_min + l)
# Show how l controls personalization strength
print("Effect of smoothing parameter l on f_u:")
print("-" * 50)
print(f"{'l':<10} {'User A f_u':<15} {'User B f_u':<15} {'Difference':<10}")
print("-" * 50)

H_a = compute_entropy(user_a_genres)  # ~0.5
H_b = compute_entropy(user_b_genres)  # ~1.609

# For this example, assume population-wide bounds
H_min, H_max = 0.5, 1.61  # You'd compute these from your user base
for l in [0, 0.25, 0.5, 1.0, 2.0]:
    f_a = compute_f_u(H_a, H_min, H_max, l)
    f_b = compute_f_u(H_b, H_min, H_max, l)
    print(f"{l:<10} {f_a:<15.3f} {f_b:<15.3f} {f_b - f_a:<10.3f}")
Effect of smoothing parameter l on f_u:
--------------------------------------------------
l          User A f_u      User B f_u      Difference
--------------------------------------------------
0          0.000           0.999           0.999     
0.25       0.184           1.000           0.815     
0.5        0.311           1.000           0.689     
1.0        0.474           1.000           0.526     
2.0        0.643           1.000           0.357

As l increases, the gap between User A and User B shrinks. At l=0, we get maximum personalization (f_u ranges from 0 to 1). At l=2, both users are clustered around 0.5, meaning they get nearly the same diversity treatment. This lets you tune how aggressively the system personalizes: start conservative with higher l, then decrease it as you gain confidence in the entropy signal.

Now let’s bring in our DPP functions from Part 2.

Code 
def greedy_map_dpp(L, k):
    N = L.shape[0]
    selected = []
    remaining = set(range(N))
    
    for _ in range(k):
        best_item, best_det = None, -1
        
        for item in remaining:
            candidate_set = selected + [item]
            det_val = np.linalg.det(L[np.ix_(candidate_set, candidate_set)])
            
            if det_val > best_det:
                
                best_det = det_val
                best_item = item
        
        selected.append(best_item)
        remaining.remove(best_item)
    
    return selected

def youtube_dpp_kernel(relevance, embeddings, alpha=0.5, sigma=1.0):
    N = len(relevance)
    L = np.zeros((N, N))
    
    # Compute pairwise distances
    # D_ij = ||embedding_i - embedding_j||²
    for i in range(N):
        for j in range(N):
            if i == j:
                # Diagonal: quality squared
                
                L[i, i] = relevance[i] ** 2
            else:
                # Off-diagonal: scaled similarity with RBF kernel
                D_ij = np.sum((embeddings[i] - embeddings[j]) ** 2)
                similarity = np.exp(-D_ij / (2 * sigma**2))
                L[i, j] = alpha * relevance[i] * relevance[j] * similarity
    
    return L

def youtube_dpp_ranking(relevance, embeddings, k, alpha=0.5, sigma=1.0):
    W = list(range(len(relevance)))  # Remaining candidate indices
    R = []  # Final ranked list
    
    while len(W) > 0:
        # Build kernel for current candidate pool
        rel_W = relevance[W]
        emb_W = embeddings[W]
        L = youtube_dpp_kernel(rel_W, emb_W, alpha=alpha, sigma=sigma)
        
        # Select up to k items from current pool
        window_size = min(k, len(W))
        M = greedy_map_dpp(L, k=window_size)
        
        # Map back to original indices
        selected_items = [W[i] for i in M]
        R.extend(selected_items)
        
        # Remove selected items from candidate pool
        W = [w for w in W if w not in selected_items]
    
    return R

Let’s create the personalized version.

def compute_user_alpha(genre_history, H_min, H_max, alpha_0, l=0):
    """Compute personalized alpha for a user based on their genre history."""
    H_u = compute_entropy(genre_history)
    f_u = compute_f_u(H_u, H_min, H_max, l)
    return f_u * alpha_0

def pdpp_ranking(relevance, embeddings, k, genre_history, H_min, H_max, alpha_0=1.0, l=0, sigma=1.0):
    """Personalized DPP ranking using user's genre history."""
    alpha_u = compute_user_alpha(genre_history, H_min, H_max, alpha_0, l)
    print(f"User alpha: {alpha_u:.4f} (base alpha_0={alpha_0})")
    return youtube_dpp_ranking(relevance, embeddings, k, alpha=alpha_u, sigma=sigma)

# 5 videos with relevance scores and genre embeddings (2D for simplicity)
relevance = np.array([0.9, 0.85, 0.8, 0.7, 0.6, 0.5])
genre = ['Action', 'Action', 'Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary']

# Embeddings: action movies are close together, others spread out
embeddings = np.array([
    [1.0, 0.0, 0.0],   # video 0: pure action
    [0.9, 0.1, 0.0],   # video 1: action (very similar to 0)
    [0.8, 0.1, 0.1],   # video 2: action (also very similar to 0)
    [0.0, 1.0, 0.0],   # video 3: comedy
    [0.0, 0.0, 1.0],   # video 4: documentary
    [0.1, 0.5, 0.5],   # video 5: mixed
])

# Global entropy bounds (you'd compute these from your user population)
H_min, H_max = 0.5, 1.61

# Compare rankings for User A (narrow taste) vs User B (diverse taste)
print("=== User A (narrow taste - prefers action) ===")
ranking_a = pdpp_ranking(relevance, embeddings, k=6, genre_history=user_a_genres, H_min=H_min, H_max=H_max, l=0)
print(f"Top 6: {[genre[i]+'('+str(relevance[i])+')' for i in ranking_a[:6]]}")

print("=== User B (diverse taste) ===")
ranking_b = pdpp_ranking(relevance, embeddings, k=6, genre_history=user_b_genres, H_min=H_min, H_max=H_max, l=0)
print(f"Top 6: {[genre[i]+'('+str(relevance[i])+')' for i in ranking_b[:6]]}")
=== User A (narrow taste - prefers action) ===
User alpha: 0.0004 (base alpha_0=1.0)
Top 6: ['Action(0.9)', 'Action(0.85)', 'Action(0.8)', 'Comedy(0.7)', 'Documentary(0.6)', 'Mixed comedy/documentary(0.5)']
=== User B (diverse taste) ===
User alpha: 0.9995 (base alpha_0=1.0)
Top 6: ['Action(0.9)', 'Comedy(0.7)', 'Documentary(0.6)', 'Mixed comedy/documentary(0.5)', 'Action(0.8)', 'Action(0.85)']

The results show exactly what we’d expect:

User A (α ≈ 0.0) gets the three action movies upfront, ranked purely by relevance. With α near zero, diversity has almost no influence. The algorithm respects their focused history: “You’ve watched action 80% of the time, so here’s the best action content first.”

User B (α ≈ 1.0) sees a completely different pattern. They get one action movie, then the algorithm immediately jumps to comedy, documentary, and mixed content. The other two action movies get pushed to positions 5 and 6. Their high entropy history translates to high α, which penalizes clustering similar items together.

Same 6 videos. Same relevance scores. Completely different experiences.

This is the power of pDPP: rather than treating diversity as a global system setting, it becomes a per-user preference learned directly from behavior. User A’s narrow feed isn’t a bug; it’s what their history tells us they want. User B’s varied feed isn’t random; it matches their demonstrated exploration pattern.

Figure: pDPP example. Credit: NotebookLM

A Production Variant: Bounded Additive Formulation

The multiplicative formula α_u = f_u × α_0 has a problem: it can be too aggressive. Look at User A’s result above: with f_u ≈ 0, they got α ≈ 0, meaning zero diversity consideration. In production, you probably don’t want to completely eliminate diversity for anyone.

The multiplicative approach also lacks intuitive bounds. If your team decides “we never want α below 0.4 or above 0.8,” the multiplicative formula doesn’t give you that control directly.

The additive bounded formulation:

Instead of multiplying, we use an additive offset centered around 0.5:

α_u = α_0 + (f_u - 0.5) × α_range

Where:

  • α_0 is your baseline (center point)
  • α_range controls how far users can deviate from baseline
  • f_u - 0.5 shifts the range to [-0.5, +0.5], so the adjustment is symmetric

This gives you explicit floor and ceiling:

  • Floor: α_0 - 0.5 × α_range
  • Ceiling: α_0 + 0.5 × α_range

For example, with α_0 = 0.6 and α_range = 0.4:

  • Minimum α = 0.6 - 0.2 = 0.4 (most focused users)
  • Maximum α = 0.6 + 0.2 = 0.8 (most diverse users)
def compute_user_alpha_bounded(genre_history, H_min, H_max, alpha_0=0.6, alpha_range=0.4, l=0):
    """Compute personalized alpha with explicit floor/ceiling bounds."""
    H_u = compute_entropy(genre_history)
    f_u = compute_f_u(H_u, H_min, H_max, l)

    # Additive formulation: center at alpha_0, vary by alpha_range
    alpha_u = alpha_0 + (f_u - 0.5) * alpha_range
    return alpha_u

def pdpp_ranking_bounded(relevance, embeddings, k, genre_history, H_min, H_max, 
                         alpha_0=0.6, alpha_range=0.4, l=0, sigma=1.0):
    """Personalized DPP with bounded additive alpha."""
    alpha_u = compute_user_alpha_bounded(genre_history, H_min, H_max, alpha_0, alpha_range, l)
    return youtube_dpp_ranking(relevance, embeddings, k, alpha=alpha_u, sigma=sigma)

print("=== User A (narrow taste) - Bounded ===")
ranking_a_bounded = pdpp_ranking_bounded(
    relevance, embeddings, k=6, genre_history=user_a_genres,
    H_min=H_min, H_max=H_max, alpha_0=0.6, alpha_range=0.4
)
print(f"Top 6: {[genre[i]+'('+str(relevance[i])+')' for i in ranking_a_bounded]}\n")

print("=== User B (diverse taste) - Bounded ===")
ranking_b_bounded = pdpp_ranking_bounded(
    relevance, embeddings, k=6, genre_history=user_b_genres,
    H_min=H_min, H_max=H_max, alpha_0=0.6, alpha_range=0.4
)
print(f"Top 6: {[genre[i]+'('+str(relevance[i])+')' for i in ranking_b_bounded]}")
=== User A (narrow taste) - Bounded ===
Top 6: ['Action(0.9)', 'Action(0.85)', 'Action(0.8)', 'Comedy(0.7)', 'Documentary(0.6)', 'Mixed comedy/documentary(0.5)']

=== User B (diverse taste) - Bounded ===
Top 6: ['Action(0.9)', 'Comedy(0.7)', 'Documentary(0.6)', 'Action(0.85)', 'Action(0.8)', 'Mixed comedy/documentary(0.5)']

Now both users stay within the guardrails we set (0.4 to 0.8):

User A (α = 0.4) still gets all three action movies first, but now there’s some diversity consideration baked in. They’re at the floor, not at zero. If we had more videos in the candidate pool, this baseline diversity would prevent completely homogeneous results.

User B (α = 0.8) gets a diversified feed, but not as extreme as before. Compare to the multiplicative version where they had α ≈ 1.0: now the second action movie appears at position 4 instead of position 5. The ceiling prevents over-diversification.

Why this matters in production:

  • You can guarantee minimum diversity for all users (prevents filter bubbles)
  • You can cap maximum diversity (prevents feeds that feel random)
  • The parameters are intuitive to tune: “baseline 0.6, users can vary ±0.2”
  • Product and engineering can align on explicit bounds before launch

The tradeoff is less dramatic personalization. The multiplicative approach gave a 10x difference between User A and B (0.0 vs 1.0). The bounded approach gives a 2x difference (0.4 vs 0.8). Which is right depends on your product goals and risk tolerance.

Production Considerations

Deploying pDPP requires a few decisions that don’t show up in toy examples.

Computing H_min and H_max

You need population-wide entropy bounds. Two approaches:

  1. Batch computation: Run a daily/weekly job over all active users, compute their entropy, store the min/max. Simple, but lags behind distribution shifts.

  2. Approximate with theory: If you have n genres, the theoretical bounds are H_min = 0 (user watches one genre) and H_max = ln(n) (uniform distribution). This avoids the batch job but may not reflect your actual user distribution.

In practice, I’d recommend starting with theoretical bounds, then validating against your actual user distribution. If 99% of users fall between 0.3 and 1.2 but your theoretical max is 2.3, you’re wasting personalization range.

Cold Start

New users have no history, so no entropy. Options:

  • Default to α_0 (the baseline) until they have N interactions
  • Use cohort-level entropy (users who signed up from similar sources tend to have similar diversity preferences)
  • Start slightly above baseline to encourage exploration, then let it settle

How Often to Update f_u

Entropy is relatively stable for established users. Someone who’s watched 500 videos won’t see their entropy swing wildly from one more view. Options:

  • Batch (daily/weekly): Simplest. Compute entropy offline, store in user profile.
  • Nearline (per session): Recompute when user starts a session. More responsive but adds latency.
  • Streaming: Update incrementally with each interaction. Most responsive but adds complexity.

For most systems, daily batch is sufficient. Entropy changes slowly.

What Counts as “Genre”?

The paper uses genre, but you can use any categorical attribute:

  • Content category (comedy, drama, news)
  • Creator/channel
  • Content length bucket (shorts vs. long-form)
  • Engagement type (videos you comment on vs. passive watches)

You can even combine multiple: compute entropy over (genre, creator_type) tuples. The key is choosing attributes that meaningfully capture “diversity” for your product.

Key Takeaways

  1. Standard DPP’s single α ignores that users have different diversity appetites. Shannon entropy over interaction history gives us a principled way to measure each user’s diversity propensity from their behavior.

  2. pDPP personalizes the relevance-diversity tradeoff per user. High-entropy users (explorers) get more diversity; low-entropy users (focused) get more relevance.

  3. The bounded additive formulation is production-friendly. It gives explicit floor/ceiling control and prevents edge cases where diversity is completely eliminated.

References & Further Reading

Previous posts in this series:

I hope this series has been useful for thinking about diversity in your recommendation systems! If you’re experimenting with pDPP or have questions, connect with me on LinkedIn to share your experience.

]]> Recommender System Machine Learning Diversity https://aayushmnit.com/posts/2026-01-01-DiversityPDppPart3/DiversityPDppPart3.html Thu, 01 Jan 2026 08:00:00 GMT Diversity in Recommendations - Determinantal Point Processes (DPP) Aayush Agrawal https://aayushmnit.com/posts/2025-12-28-DiversityDPPPart2/DiversityDPPPart2.html

Part 2 of the diversity series: Moving from greedy MMR to principled probabilistic frameworks with DPPs

Figure: DPP infographic. Credit: NotebookLM
WarningDisclaimer

The views expressed in this blog are solely my own and are not affiliated with my employer in any way. In writing this, I have not utilized any proprietary or confidential information.

In previous blog, we explored MMR as a practical way to inject diversity into recommendation results. MMR works well, but it has a fundamental limitation: it’s greedy and myopic. At each step, it only considers similarity to items already selected, not the global structure of the final set.

What if we could evaluate the diversity of an entire set at once? That’s exactly what Determinantal Point Processes (DPPs) offer. DPPs come from probability theory and give us a principled way to assign higher probability to sets of items that are both high-quality and diverse. For example, MMR might select two sci-fi thrillers that individually seem distinct from other picks, but together make the final set feel repetitive. DPPs evaluate the whole set at once, catching this overlap.

In this post, we’ll:

  • Build geometric intuition for DPPs through determinants as volume
  • Understand the quality-diversity decomposition for production systems
  • Implement greedy MAP inference step-by-step
  • Explore YouTube’s production approach with tunable parameters (α, σ)
  • Walk through the windowed algorithm for ranking entire feeds
  • Compare DPP and MMR on our video recommendation example

Linear Algebra Refresher: Determinants as Volume

Figure: Mathematical Intuition behind DPP. Credit: NotebookLM

Before diving into DPPs, we need to refresh one key concept: determinants measure volume.

Consider two vectors in 2D space. The determinant of the matrix formed by these vectors equals the area of the parallelogram they span.

Code 
import numpy as np
import matplotlib.pyplot as plt
# Two similar vectors (small area)
v1 = np.array([1, 0])
v2 = np.array([0.9, 0.1])
similar_matrix = np.array([v1, v2])
print(f"Similar vectors determinant: {np.linalg.det(similar_matrix):.3f}")

# Two orthogonal vectors (maximum area)
v1 = np.array([1, 0])
v2 = np.array([0, 1])
diverse_matrix = np.array([v1, v2])
print(f"Orthogonal vectors determinant: {np.linalg.det(diverse_matrix):.3f}")
Similar vectors determinant: 0.100
Orthogonal vectors determinant: 1.000
Code 
fig, axes = plt.subplots(1, 2, figsize=(10, 4))

def plot_parallelogram(ax, v1, v2, title, det_val):
    v1, v2 = np.array(v1), np.array(v2)
    # Plot parallelogram
    vertices = np.array([[0, 0], v1, v1 + v2, v2, [0, 0]])
    ax.fill(vertices[:, 0], vertices[:, 1], alpha=0.3, color='steelblue')
    ax.plot(vertices[:, 0], vertices[:, 1], 'b-', linewidth=2)
    
    # Plot vectors as arrows
    ax.annotate('', xy=v1, xytext=(0, 0), arrowprops=dict(arrowstyle='->', color='red', lw=2))
    ax.annotate('', xy=v2, xytext=(0, 0), arrowprops=dict(arrowstyle='->', color='green', lw=2))
    
    ax.set_xlim(-0.2, 2.0)
    ax.set_ylim(-0.2, 1.2)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_title(f"{title}\ndet = {det_val:.2f} (area)")
    ax.set_xlabel('x')
    ax.set_ylabel('y')

# Similar vectors
plot_parallelogram(axes[0], [1, 0], [0.9, 0.1], "Similar Vectors", 0.1)

# Orthogonal vectors  
plot_parallelogram(axes[1], [1, 0], [0, 1], "Orthogonal Vectors", 1.0)

plt.tight_layout()
plt.savefig('determinant_volume.png', dpi=150, bbox_inches='tight')
plt.show()

This is the geometric intuition behind DPPs: diverse sets span more volume, giving them higher determinants. When vectors point in similar directions, they span a thin, squashed parallelogram. When they’re spread out, the volume is large.

This extends to higher dimensions. For k vectors in d-dimensional space, the determinant of their Gram matrix (L = VVᵀ) measures the squared volume of the k-dimensional parallelotope they span.

Note: Gram matrix is also called the “kernel matrix” in DPP literature.

# 3 items in 3D feature space
items = np.array([
    [1, 0, 0],    # Item A
    [0, 1, 0],    # Item B  
    [0, 0, 1],    # Item C (diverse from A, B)
])
L = items @ items.T  # Gram matrix
print(f"Diverse set det(L): {np.linalg.det(L):.3f}")

# Replace C with something similar to A
items_similar = np.array([
    [1, 0, 0],      # Item A
    [0, 1, 0],      # Item B
    [0.8, 0.1, 0.1],  # Item C' (similar to A)
])
L_similar = items_similar @ items_similar.T
print(f"Similar set det(L): {np.linalg.det(L_similar):.3f}")
Diverse set det(L): 1.000
Similar set det(L): 0.010

The diverse set has 100x higher determinant. This is the core of DPPs: recommendation probability proportional to determinant means diverse sets are exponentially more likely.

DPP Fundamentals: From Volume to Probability

Now we can define what a DPP actually is. Given a ground set of N items (our candidate pool), a DPP defines a probability distribution over all possible subsets. The key formula is:

P(Y) ∝ det(L_Y)

where:

  • Y is a subset of items we might recommend
  • L is an N×N positive semi-definite kernel matrix capturing item relationships
  • L_Y is the submatrix of L indexed by items in Y

The probability of selecting a set is proportional to the determinant of its kernel submatrix. From our linear algebra refresher, we know determinants measure volume, so diverse sets that span more volume get higher probability.

Let’s use the same video recommendation example from Part 1:

relevance = np.array([0.9, 0.85, 0.8, 0.7, 0.6, 0.5])
genre = ['Action', 'Action', 'Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary']

embeddings = np.array([
    [1.0, 0.0, 0.0],   # video 0: pure action
    [0.9, 0.1, 0.0],   # video 1: action (very similar to 0)
    [0.8, 0.1, 0.1],   # video 2: action (also very similar to 0)
    [0.0, 1.0, 0.0],   # video 3: comedy
    [0.0, 0.0, 1.0],   # video 4: documentary
    [0.1, 0.5, 0.5],   # video 5: mixed
])

# Build kernel matrix L = VVᵀ (Gram matrix from embeddings)
L = embeddings @ embeddings.T
print("Kernel matrix L:")
print(np.round(L, 2))

print(f"\nSimilarity of video 0 ({genre[0]}), with other videos")
for i in range(1,len(genre)):
    print(f"Video {i}: {L[0][i]} ({genre[i]})")
Kernel matrix L:
[[1.   0.9  0.8  0.   0.   0.1 ]
 [0.9  0.82 0.73 0.1  0.   0.14]
 [0.8  0.73 0.66 0.1  0.1  0.18]
 [0.   0.1  0.1  1.   0.   0.5 ]
 [0.   0.   0.1  0.   1.   0.5 ]
 [0.1  0.14 0.18 0.5  0.5  0.51]]

Similarity of video 0 (Action), with other videos
Video 1: 0.9 (Action)
Video 2: 0.8 (Action)
Video 3: 0.0 (Comedy)
Video 4: 0.0 (Documentary)
Video 5: 0.1 (Mixed comedy/documentary)

Notice how videos 0, 1, 2 (all Action) have high similarity (~0.8-0.9), while Action vs Comedy is 0. Now let’s compare the DPP probability of two different 3-video sets:

def subset_det(L, indices):
    """Compute det(L_Y) for a subset of indices"""
    L_Y = L[np.ix_(indices, indices)]
    return np.linalg.det(L_Y)

# Set 1: Three action videos (redundant)
redundant = [0, 1, 2]  # Action, Action, Action
det_redundant = subset_det(L, redundant)
print(f"Redundant set: {[genre[i] for i in redundant]}")
print(f"  det(L_Y) = {det_redundant:.4f}")

# Set 2: Action + Comedy + Documentary (diverse)
diverse = [0, 3, 4]
det_diverse = subset_det(L, diverse)
print(f"\nDiverse set: {[genre[i] for i in diverse]}")
print(f"  det(L_Y) = {det_diverse:.4f}")

print(f"\nDiverse set is {det_diverse/det_redundant:.0f}x more likely under DPP")
Redundant set: ['Action', 'Action', 'Action']
  det(L_Y) = 0.0001

Diverse set: ['Action', 'Comedy', 'Documentary']
  det(L_Y) = 1.0000

Diverse set is 10000x more likely under DPP

The diverse set spanning Action, Comedy, and Documentary is 10000x more likely than picking three similar action videos! This is the DPP magic: it naturally penalizes redundancy by evaluating the entire set at once, rather than making greedy pairwise comparisons like MMR.

NoteWhy L(Kernel Matrix) needs to be Positive Semi-Definite (PSD)

L needs to be positive semi-definite (PSD). Practically, this means:

  1. All eigenvalues are ≥ 0
  2. Determinants of all submatrices are ≥ 0 (which we need for valid probabilities!)

Why DPPs require PSD: Since P(Y) ∝ det(L_Y), we need det(L_Y) ≥ 0 for all subsets. PSD guarantees this.

Good news for practitioners: When you build L = VVᵀ (Gram matrix from embeddings), it’s automatically PSD. So as long as you construct your kernel matrix as L = VVᵀ (or use standard similarity kernels like RBF), you’re safe. You only run into trouble if you manually construct L with arbitrary values.

Quality vs Diversity Decomposition

Figure: Quality diversity tradeoff in DPP. Credit: NotebookLM

So far our kernel matrix L only captures similarity between items. But in real recommendations, we care about both quality (relevance) and diversity. A diverse set of low-quality items is useless.

DPPs handle this elegantly through the quality-diversity decomposition:

L[i,j] = q_i × q_j × S[i,j]

where:

  • q_i is the quality (relevance) score for item i
  • S[i,j] is the similarity between items i and j (from embeddings)

This decomposes naturally: diagonal entries L[i,i] = q_i² capture item quality, while off-diagonal entries capture similarity scaled by both items’ quality. High-quality similar items create larger off-diagonal terms, which reduce the determinant (penalizing redundancy among your best items).

Let’s rebuild our kernel matrix with relevance scores:

# Similarity matrix from embeddings
S = embeddings @ embeddings.T

# Quality-weighted kernel: L[i,j] = q_i * q_j * S[i,j]
q = relevance.reshape(-1, 1)  # Column vector
L_quality = (q @ q.T) * S

print("Quality-weighted kernel L:")
print(np.round(L_quality, 3))
Quality-weighted kernel L:
[[0.81  0.688 0.576 0.    0.    0.045]
 [0.688 0.592 0.496 0.06  0.    0.06 ]
 [0.576 0.496 0.422 0.056 0.048 0.072]
 [0.    0.06  0.056 0.49  0.    0.175]
 [0.    0.    0.048 0.    0.36  0.15 ]
 [0.045 0.06  0.072 0.175 0.15  0.128]]

Let’s compare the two sets again with the new kernel matrix:

redundant = [0, 1, 2]  # Top 3 by relevance, all Action
diverse = [0, 3, 4]    # Action + Comedy + Documentary

det_redundant = subset_det(L_quality, redundant)
det_diverse = subset_det(L_quality, diverse)

print(f"Redundant set (top-3 relevance): {[f'{genre[i]} ({relevance[i]})' for i in redundant]}")
print(f" det = {det_redundant:.6f}")

print(f"\nDiverse set (mixed relevance): {[f'{genre[i]} ({relevance[i]})' for i in diverse]}")
print(f" det = {det_diverse:.6f}")

print(f"\nDiverse set is {det_diverse/det_redundant:.0f}x more likely")
Redundant set (top-3 relevance): ['Action (0.9)', 'Action (0.85)', 'Action (0.8)']
 det = 0.000037

Diverse set (mixed relevance): ['Action (0.9)', 'Comedy (0.7)', 'Documentary (0.6)']
 det = 0.142884

Diverse set is 3815x more likely

Even though the diverse set includes lower-relevance items (0.7, 0.6 vs 0.85, 0.8), the diversity bonus from spanning different genres still makes it far more probable (3815x as compared to 10000x before). The quality weighting ensures we don’t pick diverse but irrelevant items.

The intuition: Think of q_i as the “length” of item i’s feature vector. The determinant measures the volume spanned by these vectors. A high-quality item contributes a longer vector, but if two long vectors point in the same direction (similar items), the parallelogram is still flat. You need both length (quality) AND spread (diversity) to maximize volume.

Code 
fig, axes = plt.subplots(1, 3, figsize=(9, 4))

# Case 1: Short vectors, diverse directions
plot_parallelogram(axes[0], [0.5, 0], [0, 0.5], 
    "Low Quality, High Diversity\nShort vectors, spread out", 0.25)

# Case 2: Long vectors, similar directions  
plot_parallelogram(axes[1], [1.0, 0], [0.9, 0.1],
    "High Quality, Low Diversity\nLong vectors, same direction", 0.10)

# Case 3: Long vectors, diverse directions (ideal)
plot_parallelogram(axes[2], [1.0, 0], [0, 1.0],
    "High Quality, High Diversity\nLong vectors, spread out", 1.00)

plt.tight_layout()
plt.savefig('quality_diversity_volume.png', dpi=150, bbox_inches='tight')
plt.show()

This maps directly to production recsys:

  • q_i = your relevance model output (e.g., predicted click probability)
  • S[i,j] = Similarity of content embeddings (or 1 - distance in any embedding space)
  • L = the kernel that balances both signals

Greedy MAP: The Practical DPP Algorithm

We’ve established that DPPs assign higher probability to diverse, high-quality sets via determinants. The natural question: how do we find the best set?

The challenge: Finding the subset Y that maximizes det(L_Y) is called MAP (Maximum A Posteriori) inference. With N candidates, there are 2^N possible subsets. This is NP-hard in general.

The solution: Greedy MAP approximation. Build the set one item at a time, always adding whichever item increases the determinant most. This runs in O(Nk³) time and provides provably good approximations, where N is the number of candidates to select from and k is the number of items to select.

Walking Through Greedy MAP Step-by-Step

Let’s trace through selecting 3 videos from our candidate pool. At each step, we try adding every remaining candidate and pick whoever maximizes det(L_Y).

Step 1: First Selection

With an empty set, we’re just comparing single-item determinants. For a 1×1 matrix, det([L[i,i]]) = L[i,i] = q_i². So the first pick is simply the highest-quality item:

print("Step 1: Evaluating single-item determinants (= quality squared)")
for i in range(len(relevance)):
    det_single = L_quality[i, i]
    print(f"  Video {i} ({genre[i]}): det = {det_single:.3f} (relevance² = {relevance[i]**2:.3f})")

first_pick = int(np.argmax([L_quality[i,i] for i in range(len(relevance))]))
print(f"\n→ Select video {first_pick} ({genre[first_pick]})")
Step 1: Evaluating single-item determinants (= quality squared)
  Video 0 (Action): det = 0.810 (relevance² = 0.810)
  Video 1 (Action): det = 0.592 (relevance² = 0.722)
  Video 2 (Action): det = 0.422 (relevance² = 0.640)
  Video 3 (Comedy): det = 0.490 (relevance² = 0.490)
  Video 4 (Documentary): det = 0.360 (relevance² = 0.360)
  Video 5 (Mixed comedy/documentary): det = 0.128 (relevance² = 0.250)

→ Select video 0 (Action)

Step 2: Second Selection

Now it gets interesting. We evaluate each remaining candidate by computing det(L_Y) for the 2-item set:

selected = [first_pick]
print(f"Step 2: Current set = [Video {first_pick}]")
print("Evaluating candidates:")

for candidate in range(len(relevance)):
    if candidate in selected:
        continue
    
    subset = selected + [candidate]
    L_subset = L_quality[np.ix_(subset, subset)]
    det_val = np.linalg.det(L_subset)
    
    # Show the 2x2 matrix being evaluated
    print(f"\n  Video {candidate} ({genre[candidate]}):")
    print(f"    L_subset = [[{L_subset[0,0]:.3f}, {L_subset[0,1]:.3f}],")
    print(f"                [{L_subset[1,0]:.3f}, {L_subset[1,1]:.3f}]]")
    print(f"    det = {L_subset[0,0]:.3f} × {L_subset[1,1]:.3f} - {L_subset[0,1]:.3f} × {L_subset[1,0]:.3f} = {det_val:.4f}")
Step 2: Current set = [Video 0]
Evaluating candidates:

  Video 1 (Action):
    L_subset = [[0.810, 0.689],
                [0.689, 0.592]]
    det = 0.810 × 0.592 - 0.689 × 0.689 = 0.0059

  Video 2 (Action):
    L_subset = [[0.810, 0.576],
                [0.576, 0.422]]
    det = 0.810 × 0.422 - 0.576 × 0.576 = 0.0104

  Video 3 (Comedy):
    L_subset = [[0.810, 0.000],
                [0.000, 0.490]]
    det = 0.810 × 0.490 - 0.000 × 0.000 = 0.3969

  Video 4 (Documentary):
    L_subset = [[0.810, 0.000],
                [0.000, 0.360]]
    det = 0.810 × 0.360 - 0.000 × 0.000 = 0.2916

  Video 5 (Mixed comedy/documentary):
    L_subset = [[0.810, 0.045],
                [0.045, 0.128]]
    det = 0.810 × 0.128 - 0.045 × 0.045 = 0.1013

Notice how adding another Action video (high off-diagonal similarity) leads to a lower determinant, while Comedy or Documentary (zero similarity to Action) leads to a higher determinant. In this case we will pick Video 3 as it has the highest determinant of 0.3969.

# Find second pick
best_det, second_pick = -1, None
for candidate in range(len(relevance)):
    if candidate in selected:
        continue
    subset = selected + [candidate]
    det_val = np.linalg.det(L_quality[np.ix_(subset, subset)])
    if det_val > best_det:
        best_det, second_pick = det_val, candidate

print(f"\n→ Select video {second_pick} ({genre[second_pick]}), det = {best_det:.4f}")
selected.append(second_pick)

→ Select video 3 (Comedy), det = 0.3969

Step 3: Third Selection

Same process with a 3×3 determinant:

print(f"Step 3: Current set = {selected}")
print("Evaluating candidates:")

for candidate in range(len(relevance)):
    if candidate in selected:
        continue
    
    subset = selected + [candidate]
    det_val = np.linalg.det(L_quality[np.ix_(subset, subset)])
    print(f"  Video {candidate} ({genre[candidate]}): det = {det_val:.4f}")

# Find third pick
best_det, third_pick = -1, None
for candidate in range(len(relevance)):
    if candidate in selected:
        continue
    subset = selected + [candidate]
    det_val = np.linalg.det(L_quality[np.ix_(subset, subset)])
    if det_val > best_det:
        best_det, third_pick = det_val, candidate

selected.append(third_pick)
print(f"\n→ Select video {third_pick} ({genre[third_pick]})")
print(f"\nFinal selection: {[genre[i] for i in selected]}")
Step 3: Current set = [0, 3]
Evaluating candidates:
  Video 1 (Action): det = 0.0000
  Video 2 (Action): det = 0.0025
  Video 4 (Documentary): det = 0.1429
  Video 5 (Mixed comedy/documentary): det = 0.0248

→ Select video 4 (Documentary)

Final selection: ['Action', 'Comedy', 'Documentary']

The Complete Implementation

Now we can wrap this into a clean function:

def greedy_map_dpp(L, k):
    N = L.shape[0]
    selected = []
    remaining = set(range(N))
    
    for _ in range(k):
        best_item, best_det = None, -1
        
        for item in remaining:
            candidate_set = selected + [item]
            det_val = np.linalg.det(L[np.ix_(candidate_set, candidate_set)])
            
            if det_val > best_det:
                
                best_det = det_val
                best_item = item
        
        selected.append(best_item)
        remaining.remove(best_item)
    
    return selected

# Run it
result = greedy_map_dpp(L_quality, k=3)
print(f"Greedy MAP result: {[f'{genre[i]} ({relevance[i]})' for i in result]}")
Greedy MAP result: ['Action (0.9)', 'Comedy (0.7)', 'Documentary (0.6)']

Pure relevance ranking would give us Action, Action, Action. DPP gives us Action, Comedy, Documentary. Three genres covered in three slots.

Tuning the Quality-Diversity Tradeoff

Figure: Youtube’s DPP Implementation. Credit: NotebookLM

In production, you need a way to control how much the system emphasizes relevance versus diversity. YouTube’s approach, described in their 2018 paper on DPPs for recommendations, provides a clean parameterization that’s proven to work at scale.

YouTube’s Kernel Parameterization

Figure: Snapshot of DPP kenel matrix formulation from YT DPP paper

Instead of the simple L[i,j] = q_i × q_j × S[i,j] formulation, YouTube uses:

L[i,i] = q_i² (diagonal: quality squared)

L[i,j] = α × q_i × q_j × exp(-D_ij / 2σ²) for i ≠ j (off-diagonal: scaled similarity)

where:

  • q_i is the relevance score for item i
  • D_ij is the distance between items i and j in embedding space
  • α controls the overall strength of diversity (0 ≤ α ≤ 1)
  • σ controls the length scale of similarity

The key insight: The diagonal entries (quality) remain unchanged, but off-diagonal entries (similarity) are scaled by α. This gives you direct control over the diversity penalty.

Note

Note: This formulation is identical to a standard (Gaussian) radial basis function (RBF) kernel when α = 1.

Understanding the Parameters

α (diversity strength):

  • α = 1: Full diversity penalty. Similar items strongly repel each other.
  • α = 0: No diversity penalty. Reduces to pure relevance ranking.
  • α ∈ (0,1): Smooth tradeoff between quality and diversity.

σ (similarity scale):

  • Small σ: Only very similar items are considered redundant (tight clusters).
  • Large σ: Even moderately different items are penalized (broad diversity).

Let’s implement this new kernel building function:

def youtube_dpp_kernel(relevance, embeddings, alpha=0.5, sigma=1.0):
    N = len(relevance)
    L = np.zeros((N, N))
    
    # Compute pairwise distances
    # D_ij = ||embedding_i - embedding_j||²
    for i in range(N):
        for j in range(N):
            if i == j:
                # Diagonal: quality squared
                
                L[i, i] = relevance[i] ** 2
            else:
                # Off-diagonal: scaled similarity with RBF kernel
                D_ij = np.sum((embeddings[i] - embeddings[j]) ** 2)
                similarity = np.exp(-D_ij / (2 * sigma**2))
                L[i, j] = alpha * relevance[i] * relevance[j] * similarity
    
    return L

# Build kernel with YouTube's approach
L_youtube = youtube_dpp_kernel(relevance, embeddings, alpha=0.5, sigma=1.0)

print("YouTube DPP kernel (α=0.5, σ=1.0):")
print(np.round(L_youtube, 3))
YouTube DPP kernel (α=0.5, σ=1.0):
[[0.81  0.379 0.349 0.116 0.099 0.117]
 [0.379 0.722 0.337 0.132 0.103 0.126]
 [0.349 0.337 0.64  0.135 0.116 0.133]
 [0.116 0.132 0.135 0.49  0.077 0.136]
 [0.099 0.103 0.116 0.077 0.36  0.116]
 [0.117 0.126 0.133 0.136 0.116 0.25 ]]

Sweeping the α Parameter

Let’s see how α affects the diversity-quality tradeoff:

Code 
print("\nα parameter sweep (σ=1.0 fixed):")
print(f"{'α':<6} {'Selected Items':<50} {'Genres':<40}")
print("-" * 95)

alphas = [0.0, 0.25, 0.5, 0.75, 1.0]

for alpha in alphas:
    L_alpha = youtube_dpp_kernel(relevance, embeddings, alpha=alpha, sigma=1.0)
    selected = greedy_map_dpp(L_alpha, k=3)
    items_str = ", ".join([f"V{i}({relevance[i]:.2f})" for i in selected])
    genres_str = ", ".join([genre[i] for i in selected])
    print(f"{alpha:<6.2f} {items_str:<50} {genres_str:<40}")

α parameter sweep (σ=1.0 fixed):
α      Selected Items                                     Genres                                  
-----------------------------------------------------------------------------------------------
0.00   V0(0.90), V1(0.85), V2(0.80)                       Action, Action, Action                  
0.25   V0(0.90), V1(0.85), V2(0.80)                       Action, Action, Action                  
0.50   V0(0.90), V1(0.85), V3(0.70)                       Action, Action, Comedy                  
0.75   V0(0.90), V3(0.70), V4(0.60)                       Action, Comedy, Documentary             
1.00   V0(0.90), V3(0.70), V4(0.60)                       Action, Comedy, Documentary

Key observations:

  • α = 0.0–0.25: Pure relevance ranking. Off-diagonal terms are too weak to overcome quality differences, so we get the top-3 by relevance (all Action).
  • α = 0.5: Moderate diversity kicks in. The algorithm trades Video 2 (Action, 0.80) for Video 3 (Comedy, 0.70), accepting a small relevance drop for genre diversity.
  • α = 0.75–1.0: Strong diversity preference. The algorithm maximizes spread across genres (Action, Comedy, Documentary), dropping Video 1 (Action, 0.85) in favor of lower-relevance items from different genres.

The transition happens around α = 0.5, where the diversity penalty becomes strong enough to outweigh small relevance differences. In production, you’d tune α based on metrics like session diversity, user engagement, or A/B test results.

Sweeping the σ Parameter

The σ parameter controls what counts as “similar”. Let’s fix α = 0.75 and vary σ:

Code 
print("\nσ parameter sweep (α=0.75 fixed):")
print(f"{'σ':<6} {'Selected Items':<50} {'Genres':<40}")
print("-" * 95)

sigmas = [0.1, 0.5, 1.0, 2.0, 5.0]

for sigma in sigmas:
    L_sigma = youtube_dpp_kernel(relevance, embeddings, alpha=0.75, sigma=sigma)
    selected = greedy_map_dpp(L_sigma, k=3)
    items_str = ", ".join([f"V{i}({relevance[i]:.2f})" for i in selected])
    genres_str = ", ".join([genre[i] for i in selected])
    print(f"{sigma:<6.1f} {items_str:<50} {genres_str:<40}")

σ parameter sweep (α=0.75 fixed):
σ      Selected Items                                     Genres                                  
-----------------------------------------------------------------------------------------------
0.1    V0(0.90), V1(0.85), V2(0.80)                       Action, Action, Action                  
0.5    V0(0.90), V3(0.70), V4(0.60)                       Action, Comedy, Documentary             
1.0    V0(0.90), V3(0.70), V4(0.60)                       Action, Comedy, Documentary             
2.0    V0(0.90), V3(0.70), V1(0.85)                       Action, Comedy, Action                  
5.0    V0(0.90), V1(0.85), V2(0.80)                       Action, Action, Action

Key observations:

  • σ = 0.1: Very tight similarity threshold. Only near-identical items are penalized as redundant, so the diversity effect barely activates. Result: top-3 by relevance (all Action).
  • σ = 0.5–1.0: Moderate similarity scale. Items within the same genre are considered similar enough to penalize, leading to genre mixing (Action, Comedy, Documentary).
  • σ = 2.0: Broader similarity scale. Even moderately different items start to look similar. The algorithm becomes more willing to sacrifice relevance for diversity.
  • σ = 5.0: Very broad similarity scale. Almost everything looks similar, which actually weakens the diversity signal. The algorithm struggles to distinguish between items and reverts toward relevance-based ranking.

In practice, σ should match the natural scale of your embedding space. If your embeddings are normalized (unit length), σ ≈ 1.0 is a reasonable starting point. Tune it based on how strict you want diversity: smaller σ for subtle redundancy detection within tight clusters, larger σ for aggressive spreading across broader categories.

Visualizing the RBF Kernel

The exp(-D_ij / 2σ²) term is an RBF (Radial Basis Function) kernel. Let’s visualize how it maps distances to similarities:

Code 
fig, ax = plt.subplots(figsize=(8, 5))

distances = np.linspace(0, 3, 100)

for sigma in [0.5, 1.0, 2.0]:
    similarities = np.exp(-distances**2 / (2 * sigma**2))
    ax.plot(distances, similarities, label=f'σ={sigma}', linewidth=2)

ax.set_xlabel('Distance D_ij (embedding space)')
ax.set_ylabel('Similarity exp(-D²/2σ²)')
ax.set_title('RBF Kernel: How Distance Maps to Similarity')
ax.legend()
ax.grid(True, alpha=0.3)
ax.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
plt.savefig('rbf_kernel.png', dpi=150, bbox_inches='tight')
plt.show()

Interpretation: With σ = 1.0, items at distance 1.0 have similarity ≈ 0.6 (moderate penalty), while items at distance 2.0 have similarity ≈ 0.14 (weak penalty). Smaller σ makes the kernel “sharper” — only very close items are considered similar.

Practical Guidance

Tuning strategy:

  1. Start with α = 0.5, σ = 1.0: This is a reasonable default that balances quality and diversity.

  2. Tune α first:

    • Sweep α ∈ [0.2, 0.8] in 0.1 increments
    • Measure offline diversity metrics (intra-list distance, genre entropy) vs relevance metrics (nDCG)
    • Plot the Pareto frontier to find acceptable tradeoffs

  3. Then tune σ:

    • If items are still too similar, decrease σ (tighter similarity definition)
    • If diversity is too aggressive, increase σ (looser similarity definition)
    • σ should roughly match the typical distance scale in your embedding space

  4. A/B test top candidates: Run online experiments with 2-3 (α, σ) pairs to validate offline metrics.

Typical values in production:

  • Balanced systems: α ≈ 0.4-0.6, σ ≈ 0.5-1.5
  • Quality-focused (e.g., search): α ≈ 0.2-0.4
  • Discovery-focused (e.g., explore page): α ≈ 0.6-0.9

The beauty of YouTube’s parameterization is that it decouples quality (diagonal) from diversity (off-diagonal), giving you two independent knobs to tune the system behavior.

YouTube’s Production DPP Implementation

The greedy MAP algorithm we implemented earlier works well for selecting a fixed top-k set. However, YouTube’s production system needs to rank entire feeds of hundreds or thousands of videos. Their 2018 paper introduces a windowed approach that processes candidates in batches, making DPPs scalable for large-scale recommendation systems.

The Windowed Algorithm

Instead of selecting all k items at once from the full candidate pool, YouTube’s algorithm:

  1. Selects k items from the current candidate pool using greedy MAP
  2. Adds them to the result list R
  3. Removes selected items from the pool and updates the kernel matrix
  4. Repeats until all candidates are ranked (or enough items are selected)

This approach has several advantages:

  • Scalability: Processes N items in O(Nk²) time instead of O(Nk³)
  • Diversity within pages: Ensures each “window” of k items is diverse, not just the top-k
  • Memory efficiency: Can process large pools without computing full N×N kernels
  • Flexibility: Can stop early after ranking enough items for the user’s viewport

Here’s the algorithm from YouTube’s paper:

Figure: Snapshot of DPP algorithm from YT DPP paper 
def youtube_dpp_ranking(relevance, embeddings, k, alpha=0.5, sigma=1.0):
    W = list(range(len(relevance)))  # Remaining candidate indices
    R = []  # Final ranked list
    
    while len(W) > 0:
        # Build kernel for current candidate pool
        rel_W = relevance[W]
        emb_W = embeddings[W]
        L = youtube_dpp_kernel(rel_W, emb_W, alpha=alpha, sigma=sigma)
        
        # Select up to k items from current pool
        window_size = min(k, len(W))
        M = greedy_map_dpp(L, k=window_size)
        
        # Map back to original indices
        selected_items = [W[i] for i in M]
        R.extend(selected_items)
        
        # Remove selected items from candidate pool
        W = [w for w in W if w not in selected_items]
    
    return R

Walking Through the Windowed Algorithm

Let’s trace through ranking all 6 videos with window size k=2:

Code 
print("YouTube Windowed DPP (k=2, α=0.75, σ=1.0)")
print("=" * 70)

# We'll manually trace the first few iterations
W = list(range(len(relevance)))
R = []
k_window = 2
alpha, sigma = 0.75, 1.0

iteration = 1
while len(W) > 0 and iteration <= 3:  # Show first 3 iterations
    print(f"\nIteration {iteration}:")
    print(f"  Remaining pool W: {W}")
    print(f"  Pool genres: {[genre[i] for i in W]}")
    
    # Build kernel for current pool
    rel_W = relevance[W]
    emb_W = embeddings[W]
    L = youtube_dpp_kernel(rel_W, emb_W, alpha=alpha, sigma=sigma)
    
    # Select k items
    window_size = min(k_window, len(W))
    M = greedy_map_dpp(L, k=window_size)
    selected_items = [W[i] for i in M]
    
    print(f"  Selected from pool: {M} (pool indices)")
    print(f"  Maps to original: {selected_items}")
    print(f"  Genres: {[genre[i] for i in selected_items]}")
    print(f"  Relevance: {[relevance[i] for i in selected_items]}")
    
    R.extend(selected_items)
    W = [w for w in W if w not in selected_items]
    
    iteration += 1

print(f"\n{'─'*70}")
print(f"Final ranking R: {R}")
print(f"Genres: {[genre[i] for i in R]}")
YouTube Windowed DPP (k=2, α=0.75, σ=1.0)
======================================================================

Iteration 1:
  Remaining pool W: [0, 1, 2, 3, 4, 5]
  Pool genres: ['Action', 'Action', 'Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary']
  Selected from pool: [0, 3] (pool indices)
  Maps to original: [0, 3]
  Genres: ['Action', 'Comedy']
  Relevance: [np.float64(0.9), np.float64(0.7)]

Iteration 2:
  Remaining pool W: [1, 2, 4, 5]
  Pool genres: ['Action', 'Action', 'Documentary', 'Mixed comedy/documentary']
  Selected from pool: [0, 2] (pool indices)
  Maps to original: [1, 4]
  Genres: ['Action', 'Documentary']
  Relevance: [np.float64(0.85), np.float64(0.6)]

Iteration 3:
  Remaining pool W: [2, 5]
  Pool genres: ['Action', 'Mixed comedy/documentary']
  Selected from pool: [0, 1] (pool indices)
  Maps to original: [2, 5]
  Genres: ['Action', 'Mixed comedy/documentary']
  Relevance: [np.float64(0.8), np.float64(0.5)]

──────────────────────────────────────────────────────────────────────
Final ranking R: [0, 3, 1, 4, 2, 5]
Genres: ['Action', 'Comedy', 'Action', 'Documentary', 'Action', 'Mixed comedy/documentary']

Key observations:

  • Iteration 1: From the full pool, selects the highest-quality Action video (0.9) and highest-quality diverse option (Comedy, 0.7)
  • Iteration 2: From remaining items, selects another Action video (0.85) paired with Documentary (0.6) for diversity
  • Iteration 3: Clears the remaining pool with the last Action and Mixed genre video

Notice how diversity is enforced within each window of 2 items, not just globally. This ensures users see varied content throughout their feed, not just at the top.

Comparing Single-Pass vs Windowed

Let’s compare the two approaches side-by-side:

Code 
# Single-pass: select top-3 from full pool at once
L_full = youtube_dpp_kernel(relevance, embeddings, alpha=0.75, sigma=1.0)
single_pass = greedy_map_dpp(L_full, k=3)

# Windowed: select in windows of k=2
windowed_result = youtube_dpp_ranking(relevance, embeddings, k=2, alpha=0.75, sigma=1.0)

print("Comparison of top-3 selections:")
print(f"{'Method':<20} {'Indices':<15} {'Genres':<50} {'Relevance'}")
print("─" * 100)

single_genres = [genre[i] for i in single_pass]
single_rel = [float(relevance[i]) for i in single_pass]
print(f"{'Single-pass':<20} {str(single_pass):<15} {', '.join(single_genres):<50} {single_rel}")

windowed_top3 = windowed_result[:3]
windowed_genres = [genre[i] for i in windowed_top3]
windowed_rel = [float(relevance[i]) for i in windowed_top3]
print(f"{'Windowed (k=2)':<20} {str(windowed_top3):<15} {', '.join(windowed_genres):<50} {windowed_rel}")
Comparison of top-3 selections:
Method               Indices         Genres                                             Relevance
────────────────────────────────────────────────────────────────────────────────────────────────────
Single-pass          [0, 3, 4]       Action, Comedy, Documentary                        [0.9, 0.7, 0.6]
Windowed (k=2)       [0, 3, 1]       Action, Comedy, Action                             [0.9, 0.7, 0.85]

The windowed approach includes Video 1 (Action, 0.85) in the top-3 because it was selected in the second window after Videos 0 and 3 were removed. This demonstrates how window size k affects the ranking: smaller windows enforce local diversity but may allow similar items to appear nearby across window boundaries.

In the paper, it’s mentioned that typically the candidate pool is several hundreds of videos and sub-window size is a dozen (12) or so videos.

The windowed algorithm is YouTube’s production solution for bringing the mathematical elegance of DPPs to real-world scale. By processing items in batches, it maintains diversity guarantees while handling feeds with thousands of candidates in milliseconds.

System Design Integration

Now that we’ve covered DPP theory and implementation, let’s see how it fits into a production recommendation system. Recall from Part 1 the standard recommendation pipeline:

Figure: Typical RecSys pipeline showing where DPP fits in the recommendation pipeline

DPPs slot into the Re-ranking stage, right where we used MMR in Part 1. Here’s how the stages work together:

1. Candidate Generation (Recall)

  • Retrieve ~1000 candidates from multiple sources (collaborative filtering, content-based, trending)
  • Goal: High recall, cast a wide net
  • Output: Large pool of potentially relevant items

2. Scoring (Ranking)

  • ML model predicts relevance for each candidate (e.g., click probability, watch time)
  • Features: user history, item metadata, context
  • Output: Relevance scores q_i for each item, typically in 200-500 candidates

3. Re-ranking (Diversity)

  • This is where DPP comes in
  • Input: Scored candidates + content embeddings
  • Build kernel matrix: L[i,j] = α × q_i × q_j × exp(-D_ij / 2σ²)
  • Apply windowed DPP to produce final ranked list
  • Output: Diverse, high-quality feed. Batches of size 6-12 candidates

4. Business Logic

  • Apply final filters (content policy, already-seen, etc.)
  • Insert ads, promotions
  • Output: User-facing feed

The key advantage of DPPs is that they give you a mathematically principled way to balance quality and diversity, backed by decades of research in probability theory. Once you understand the volume-based intuition, the system becomes just as interpretable as MMR but with stronger theoretical guarantees.

Key Takeaways

  • DPPs provide a mathematically principled framework for diversity. Unlike greedy approaches like MMR that make local decisions, DPPs evaluate entire sets at once through determinants. This volume-based intuition (diverse sets span more space) translates directly into probability: P(Y) ∝ det(L_Y).

  • The quality-diversity decomposition maps naturally to production recsys. The kernel formulation L[i,j] = α × q_i × q_j × similarity(i,j) cleanly separates relevance (from your ML model) and diversity (from content embeddings). You get two interpretable knobs: α controls how much diversity matters, σ controls what counts as “similar.”

  • Greedy MAP is deterministic and faster. While exact MAP inference is NP-hard, the greedy algorithm runs in O(Nk²) time and provides strong approximations. Unlike sampling-based approaches, it’s deterministic and stable, making it suitable for A/B testing and debugging.

  • YouTube’s windowed approach scales to video feeds. By processing candidates in batches of k items (typically 6-12), the algorithm maintains diversity throughout long feeds, not just at the top. This makes DPPs practical for ranking hundreds or thousands of candidates in milliseconds.

References & Further Reading

Previous blog post:

I hope you found this deep dive into DPPs useful! If you’re implementing diversity in your recommendation system, I’d love to hear about your experience. Connect with me on LinkedIn to share your thoughts or questions.

]]> Recommender System Machine Learning Diversity https://aayushmnit.com/posts/2025-12-28-DiversityDPPPart2/DiversityDPPPart2.html Sun, 28 Dec 2025 08:00:00 GMT Diversity in Recommendations - Maximal Marginal Relevance (MMR) Aayush Agrawal https://aayushmnit.com/posts/2025-12-25-DiversityMMRPart1/DiversityMMRPart1.html

Beyond Accuracy: Using Maximal Marginal Relevance to create more diverse and engaging recommendations.

Figure 1: Diverse recommendations cover multiple user intents. Credit: Gemini
WarningDisclaimer

The views expressed in this blog are solely my own and are not affiliated with my employer in any way. In writing this, I have not utilized any proprietary or confidential information.

Picture this: you open YouTube looking for something to watch. The algorithm shows you 10 videos, and they’re all basketball highlights from the same game. Sure, you like basketball, but maybe right now you’re in the mood for that ML lecture you’ve been putting off. Or maybe something else entirely.

Most recommendation systems score each item independently based on predicted relevance. The result? If your model is confident you like basketball, you get basketball. All basketball. The system optimizes for the single most likely intent while ignoring the possibility that you might have multiple interests, or that its prediction might simply be wrong.

Diversity addresses this by explicitly penalizing redundancy. Instead of asking “what’s most relevant?”, we ask “what’s most relevant that we haven’t already covered?”

This isn’t just theoretically appealing. Diversity boosting delivers real results in production systems and can dramatically improve engagement metrics. Nobel laureate Harry Markowitz called diversification “the only free lunch in finance” when applied to portfolio management. The same principle applies to recommendations: hedge your bets across the user’s possible intents.

In this post, we’ll:

  • Explain why diversity matters and when pointwise scoring fails
  • Introduce Maximal Marginal Relevance (MMR), a simple and battle-tested diversity method
  • Walk through the algorithm step-by-step with code
  • Explore how the λ parameter lets you tune the quality-diversity tradeoff

Why Diversity Matters

Recommendation systems face a fundamental challenge: they don’t know what the user wants right now. They have historical signals, but a user who watched NBA highlights yesterday might be looking for a cooking video today. A user who binges Peppa Pig with their kids on weekends has different needs on a Tuesday night.

Pointwise ranking models score each candidate independently. If your embedding space clusters similar content together (as it should), similar items receive similar scores. The top-K results end up being variations of the same thing.

This creates several problems:

  • Missed opportunities: If the user’s current intent doesn’t match your top prediction, you’ve wasted valuable screen real estate
  • Filter bubbles: Users get stuck in narrow content loops, reducing long-term engagement
  • Poor fault tolerance: When your model is wrong about intent, it’s wrong across the entire slate

Diversity isn’t unique to recommendations. Portfolio managers diversify across asset classes to reduce risk. Autonomous driving systems like IntentNet diversify over possible agent intents rather than betting on a single predicted trajectory. The core insight is the same: when you’re uncertain about what the environment wants, spread your bets.

The business case is compelling. Users who see a mix of content that spans their interests are more likely to find something they want right now, and more likely to come back.

The Problem Setup

Let’s make this concrete with an example. Imagine you have 6 candidate videos for a user:

Code 
import numpy as np
import pandas as pd
from fastcore.utils import L
from sklearn.metrics.pairwise import cosine_similarity

relevance = np.array([0.9, 0.85, 0.8, 0.7, 0.6, 0.5])
genre = L(['Action', 'Action', 'Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary'])

df = pd.DataFrame({
    'video_id': [i for i in range(len(relevance))],
    'genre': genre,
    'relevance': relevance
})
df
video_id genre relevance
0 0 Action 0.90
1 1 Action 0.85
2 2 Action 0.80
3 3 Comedy 0.70
4 4 Documentary 0.60
5 5 Mixed comedy/documentary 0.50

Pure relevance ranking returns: Action, Action, Action, Comedy, Documentary, Mixed

The top three slots go to nearly identical content. If the user wanted comedy, they have to scroll past three action videos. If they wanted a documentary, it’s buried at position 5.

Now consider what happens with diversity-aware ranking: Action, Comedy, Documentary, Action, Mixed

The slate now covers three distinct genres in the top three positions. The second action video drops to position 4 since its marginal value is low once we’ve already shown an action video.

This is the core insight: marginal relevance matters more than absolute relevance. The 3rd action video isn’t 3x as valuable as the first. Each additional similar item has diminishing returns.

We can represent each video as an embedding vector over genres:

embeddings = np.array([
    [1.0, 0.0, 0.0],   # video 0: pure action
    [0.9, 0.1, 0.0],   # video 1: action (very similar to 0)
    [0.8, 0.2, 0.0],   # video 2: action (also very similar to 0)
    [0.0, 1.0, 0.0],   # video 3: comedy
    [0.0, 0.0, 1.0],   # video 4: documentary
    [0.1, 0.5, 0.5],   # video 5: mixed
])

similarities = cosine_similarity(embeddings)

print(f"Similarity b/w Video 0 and 1 is {similarities[0,1]:.2f}")
print(f"Similarity b/w Video 0 and 3 is {similarities[0,3]:.2f}")
Similarity b/w Video 0 and 1 is 0.99
Similarity b/w Video 0 and 3 is 0.00

Videos 0 and 1 have cosine similarity near 1.0. Videos 0 and 3 are orthogonal (similarity = 0). This similarity structure is exactly what MMR exploits to rerank results.

Maximal Marginal Relevance: The Algorithm

Figure 2: MMR infographic. Credit: NotebookLM

MMR was introduced by Carbonell and Goldstein in 1998 for document retrieval and summarization. The core formula is:

MMR = ArgMax[λ × Rel(item, user) - (1-λ) × Max(Sim(item, selected_items))]

Where:

  • item = candidate video not yet selected
  • user = user profile
  • selected_items = set of items already chosen for the slate
  • λ = trade-off parameter between relevance and diversity (0 to 1)

You can see the first part λ × Rel(item, user) is just relevance. The second part (1-λ) × Max(Sim(item, selected_items)) is the diversity penalty.

The algorithm is greedy and iterative:

  1. First pick: Select the item with highest relevance (the similarity penalty term is zero since nothing has been selected yet)
  2. Subsequent picks: For each remaining candidate, compute relevance minus a penalty for similarity to the most similar already-selected item. Pick the candidate with the highest score.
  3. Repeat until you’ve filled K slots

The key insight is the Max in the second term. You’re penalized based on your similarity to whichever selected item you’re most similar to. Being different from 2 out of 3 selected items doesn’t help if you’re a near-duplicate of the third.

The λ parameter controls the tradeoff:

  • λ = 1.0 → Pure relevance ranking, no diversity
  • λ = 0.0 → Maximum diversity, ignore relevance
  • λ = 0.5 → Equal weight to both
  • λ = 0.7 → The paper’s suggested starting point

Walking Through MMR Step-by-Step

Let’s trace through the algorithm by hand before writing any code.

Step 1: First selection

With an empty slate, there’s no diversity penalty.

selected_item = np.argmax(relevance)
print(f"First selected item: {int(selected_item)}, Relevance: {relevance[selected_item]}")
First selected item: 0, Relevance: 0.9

We pick the highest relevance item: Video 0 (Action, relevance = 0.9).

Step 2: Second selection

Now we need to choose between the remaining videos. Let’s compare Video 1 (another Action, relevance=0.85) vs Video 3 (Comedy, relevance=0.7)

lam = 0.5
video_1_mmr_score = lam * relevance[1] - (1-lam) * cosine_similarity(embeddings[[1]], embeddings[[selected_item]])
video_3_mmr_score = lam * relevance[3] - (1-lam) * cosine_similarity(embeddings[[3]], embeddings[[selected_item]])
print(f"MMR score of Video 1: {video_1_mmr_score.item():.3f}")
print(f"MMR score of Video 3: {video_3_mmr_score.item():.3f}")
MMR score of Video 1: -0.072
MMR score of Video 3: 0.350

Video 3 wins despite having lower relevance. The similarity penalty makes Video 1’s score negative. This is MMR in action: redundancy is penalized so heavily that a lower-quality diverse item beats a high-quality duplicate.

The full implementation

def mmr_select(relevance, embeddings, k=3, lam=0.5):
    n = len(relevance)
    selected = []
    remaining = list(range(n))

    # First pick: pure relevance (no penalty term)
    first = int(np.argmax(relevance))
    selected.append(first)
    remaining.remove(first)

    # Subsequent picks: balance relevance vs redundancy
    while len(selected) < k:
        best_score, best_idx = -float('inf'), None

        for i in remaining:
            # Penalty = similarity to most similar already-selected item
            max_sim = max(similarities[i, j] for j in selected)
            mmr_score = lam * relevance[i] - (1 - lam) * max_sim

            if mmr_score > best_score:
                best_score, best_idx = mmr_score, i

        selected.append(best_idx)
        remaining.remove(best_idx)

    return selected
result = mmr_select(relevance, embeddings, k=3, lam=0.5)
print(genre[result])
['Action', 'Comedy', 'Documentary']

Pure relevance ranking would give us Action, Action, Action. MMR gives us Action, Comedy, Documentary. Three genres covered in three slots.

Tuning λ: The Quality-Diversity Tradeoff

The λ parameter is your dial between relevance and diversity. Let’s see how it affects the ranking:

for lam in [0.3, 0.5, 0.7, 0.9, 1.0]:
    result = mmr_select(relevance, embeddings, k=6, lam=lam)
    print(f"λ={lam}: {genre[result]}")
λ=0.3: ['Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary', 'Action', 'Action']
λ=0.5: ['Action', 'Comedy', 'Documentary', 'Action', 'Action', 'Mixed comedy/documentary']
λ=0.7: ['Action', 'Comedy', 'Documentary', 'Action', 'Action', 'Mixed comedy/documentary']
λ=0.9: ['Action', 'Action', 'Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary']
λ=1.0: ['Action', 'Action', 'Action', 'Comedy', 'Documentary', 'Mixed comedy/documentary']

At λ=1.0 we recover pure relevance ranking. As λ decreases, diverse items get promoted earlier in the slate.

Choosing λ in practice:

  • Start with λ=0.7 (the original paper’s suggestion)
  • A/B test different values against engagement metrics

System Design of Diversity Boosting

In terms of system design, diversity boosting can be added at various stages by using different techniques, but MMR falls into the class of re-rank algorithms which are typically applied after pointwise scoring.

Figure 3: Typical RecSys pipeline showing where MMR fits in the recommendation pipeline

The diagram above shows a typical recsys pipeline :

  1. Candidate Videos → Raw pool of items to consider
  2. Pointwise video scorer → Takes user features and video features, outputs a relevance score per video independently
  3. Ranking and Diversity Boosting (highlighted in green) → This is where MMR lives. It takes the pointwise scores and reorders them to balance relevance with diversity
  4. Policy → Business rules layer (e.g., content safety filters, frequency caps)
  5. User → Final slate shown

MMR is a post-scoring reranker. The pointwise model does its job first (predicting relevance), then MMR adjusts the ordering to reduce redundancy. This is a clean separation of concerns - you don’t need to retrain your ranking model to add diversity.

Key Takeaways

  1. Pointwise ranking has a blind spot. Scoring items independently leads to redundant results when similar items cluster in your embedding space.

  2. Diversity hedges against uncertainty. When you don’t know the user’s current intent, covering multiple possibilities increases the chance of a hit.

  3. MMR is simple and effective. The algorithm is greedy, easy to implement, and has powered production systems for decades.

I hope you enjoyed reading it. If there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn.

]]> Recommender System Machine Learning Diversity https://aayushmnit.com/posts/2025-12-25-DiversityMMRPart1/DiversityMMRPart1.html Thu, 25 Dec 2025 08:00:00 GMT Why You Should Write Technical Blogs (and How to Start) Aayush Agrawal https://aayushmnit.com/posts/2025-12-13-WhyYouShouldBlog/WhyYouShouldBlog.html

A guide on why technical blogging is one of the best investments you can make in your career, along with practical tips to get started.

One of my favorite hobbies outside of work is sharing technical knowledge through blog posts. Over the years, writing has opened doors I never expected: speaking at AI meetups, opportunities to author books with famous publishers, and countless LinkedIn conversations that turned into opportunities. In this post, I want to share why I think every technical person should consider blogging, and give you some practical tips to get started.

My blogging setup

I publish on my personal website, which I host through GitHub Pages and build using Quarto. Jupyter notebooks are my tool of choice for writing since they let me mix text, code and output of that code. Quarto handles converting them into HTML. After letting a post sit on my site for a few days, I cross-post to Medium for greater visibility.

My posts focus on AI, ML, and Python development. I try to center my writing around real examples of the techniques I’m discussing so readers can gain practical knowledge they can apply to their own use cases. I typically write about:

  • Niche techniques I use at work
  • Emerging technologies (like generative AI lately)
  • Topics I’m actively studying and researching

Five reasons to start blogging

1. Increase your visibility and reach

Many people have reached out to me through LinkedIn or email asking about topics I’ve written about. Sometimes these turn into job opportunities. Blogs are like resumes, but better. They show what you actually know, not just where you’ve worked.

2. It helps you learn

Organizing knowledge forces you to synthesize your own ideas. One reliable test of whether you understand something is whether you can explain it to someone else. Writing a blog post is a great way to do that. I often discover gaps in my own understanding while trying to explain a concept clearly.

3. It improves your writing skills

The more you write, the more comfortable you become expressing your ideas. This carries over to work emails, design docs, and presentations.

4. Establish yourself as an expert

Through my writing, I’ve landed opportunities to present at local AI meetups, write books with well-known publishers, and review technical books before their release. None of this would have happened without a public body of work people could find.

5. It saves time (DRY principle)

Don’t Repeat Yourself applies to explanations too. Instead of spending two hours giving someone an overview of a topic, I can point them to an article. I also refer back to my own posts as notes, which saves me time when I need to refresh my memory on something I worked on years ago.

Tips for writing your first post

If you’re convinced but don’t know what to write about, here are some ideas that have worked for me:

  1. Curate and comment: Make a list of links to blog posts, articles, or papers you like. Write brief summaries and highlight what you found useful. This is low-pressure and helpful to others.

  2. Summarize what you learned: Attended a conference? Taking a class? Write up your notes and takeaways.

  3. Turn emails into posts: Any technical explanation you’ve written more than once in email or Slack should become a blog post.

  4. Write for yourself from a year ago: You’re best positioned to help people one step behind you. Your background, your style, and your knowledge level will give a different perspective than existing tutorials.

  5. Answer: what would have helped me six months ago?: Think about problems you struggled with and write the guide you wish existed.

  6. Don’t be a perfectionist: Post your writing sooner rather than later. Waiting just builds pressure and you end up writing less. You can always update posts later.

TipAdditional resources for motivation

Getting started

You don’t need a fancy setup. GitHub Pages is free. Quarto or Jekyll can turn markdown or notebooks into a website with minimal configuration. Medium works fine if you just want to write without managing infrastructure.

The hardest part is starting. Pick one thing you learned recently and write 500 words about it. That’s it. You can refine your process after you’ve published a few posts.

I hope you enjoyed reading this. If there is any feedback on the blog post, feel free to reach out on LinkedIn.

]]> Technical Writing Career Growth https://aayushmnit.com/posts/2025-12-13-WhyYouShouldBlog/WhyYouShouldBlog.html Sat, 13 Dec 2025 08:00:00 GMT Meta MLE Interview Preparation Guide Aayush Agrawal https://aayushmnit.com/posts/2024-12-15-MLInterviewPrep/MLInterviewPrep.html

A collection of resources while preparing for MLE interviews at Meta or other big tech companies.

WarningDisclaimer

The views expressed in this blog are solely my own and are not affiliated with my employer in any way. In writing this, I have not utilized any proprietary or confidential information. Instead, I have relied entirely on publicly available resources that were instrumental in my preparation for interviews at Meta.

Blog Resources: MLE Interview Prep Template. This resource contains a template to track your coding questions, ML system design notes, and a list of questions to prepare for behavioral interviews.

Not long ago, I transitioned from a Senior ML Scientist role at Microsoft to a Machine Learning Engineer position at Meta, and the journey was anything but quick. The preparation process was extensive, especially since it was my first experience with LeetCode-style coding interviews and ML system design interviews. While there are many resources available for preparation, I’ll be sharing the ones that helped me navigate and succeed in this challenging process.

I will try to cover the following -

  1. Interview Process Overview: Break down the entire interview process, from initial screenings to the final onsite interviews, and share tips on what to expect at each stage.

  2. Coding Interviews: Dive into the LeetCode-style coding interviews, how I prepared for them, and the strategies that worked best for me.

  3. ML System Design Interviews: Explore the ML system design interviews, offering insights into the key concepts you need to know, how to approach open-ended design problems, and the resources that helped me build a strong foundation.

  4. Behavioral Interviews: Finally, I’ll talk about the behavioral interviews, and how I prepared to effectively communicate my past work and problem-solving approach.

1 Interview Process Overview

The interview process typically starts with a screening round, followed by a more extensive onsite round.

For the phone screen, the format varies depending on the level you’re interviewing for:

  • E4/E5:
    • 5 min: Introduction
    • 35 mins: Two leetcode coding problems (Easy/Medium)
    • 5 mins: Questions for interviewer
  • E6:
    • 2 min: Introduction
    • 25 mins: One or two leetcode coding problems
    • 15 mins: Behavioral interview
    • 3 mins: Question for interviewer

Once you clear the phone screen you will be invited for on-site interviews which include coding round, ML system design round and behavioral round. Again, the composition varies depending on the level you’re interviewing for:

  • E4/E5:
    • 2 coding rounds
    • 1 ML system design
    • 1 Behavioral
  • E6
    • 2 coding rounds
    • 2 ML system design
    • 1 Behavioral

The level of preparation needed for both phone interviews and onsite rounds is similar.

While good performance in coding round is bare minimum you need to pass the loop, leveling is decided by how well you do in system design and behavioral round.

Next, we’ll dive deeper into each type of interview and explore how to best prepare for them.

2 Coding Interviews

Preparing for this was the most time-consuming part for me, especially since I don’t come from a CS background. Although I worked as an ML Scientist at Microsoft, I had limited exposure to these types of problems in the real world. However, mastering these concepts is essential for entry in high-tech software/ML roles, so it’s important to invest time in thorough preparation.

Coding round interviews unsurprisingly focus heavily on coding. A typical Interview structure looks like the following -

  • 5 min: Introduction
  • 35 mins: Two leetcode coding problems
  • 5 mins: Questions for interviewer

You are mostly expected to code on a plain text notepad with execution disabled. To get a more realistic idea of how a coding interview environment looks like watch this mock interview by interviewing.io -

Python interview with an interviewing.io engineer: Print k largest elements

2.1 A structured approach to solving coding problems in interview

When tackling a coding problem, following this structured approach can be very helpful:

  1. Ask Clarifying Questions(~3mins): When the problem is presented, read it aloud to ensure you fully understand the requirements before jumping to a solution. Ask follow-up questions to clarify any ambiguities. This might involve discussing test cases, considering edge cases, and understanding the expected input range or type. For example, think about how the solution should handle null inputs or extreme values. The ideal state is to get an alignment with your interviewer by writing out some test cases and expected output for the same.

  2. Plan Your Approach(~5 mins): Outline your solution strategy and explain it to your interviewer while typing it out in the shared text window. Break down the problem into smaller parts if possible and decide on the most appropriate algorithm or data structure and discuss any trade-offs you are making and write down potential time and space complexity of the solution you are proposing. Once your interviewer agrees with your approach and then ask permission to code it out.

  3. Write the code(~5 mins): Implement your solution, keeping your code clean and well-organized. As you code, ensure that you handle edge cases. Make sure to name your functions, classes and variables appropriately so anybody reading your code can follow.

  4. Pseudo run your solution(~2 mins): Manually run your code against various test cases while explaining it to your interviewer, including both typical and edge cases, to ensure it behaves as expected. This will help you find potential bugs and an opportunity to correct them before your interviewer points it out.

  5. Close(~2 mins): Explain time and space complexity of the solution and answer any follow up questions your interviewer might have.

By following these steps, you can effectively navigate coding problems and demonstrate a clear, methodical problem-solving approach.

TipUseful Tips for Coding Interviews

  1. Keep your introduction brief (~30 seconds) to have more time for solving the problem. For example - “Hey I am [Your Name], I currently work as [Title] at [Employer Name]. I have been working here from past [N] years. I’m now seeking new opportunities, which brings me here today.”

  2. If you’re running out of time, it’s acceptable to manually walk through one or two test cases with your interviewer. You can then suggest moving on to the next question to ensure you cover everything within the allotted time.

  3. It’s okay to ask for help from an interviewer if you are stuck on a problem

As you can see from above, coding rounds are really fast paced and you need to be well prepared to get through it. That brings us to prepration.

2.2 How to prepare for coding interviews

Here is a simple guide on how to prepare -

  1. Purchase a leetcode subscription: This website is the only paid resource you need to prepare for coding interviews.

  2. Getting started with Leetcode learn: If you are like me who doesn’t come from a CS degree then going through leetcode learn cards is a good starting point. Here is the structure I followed-

    • Array
    • Linked List
    • Stack & Queue
    • Array & Strings
    • Binary Tree
    • Binary Search
    • Binary Search Tree
    • Heap
    • Graph
    • Sorting
    • Dynamic Programming

  3. Following Neetcode.io Roadmap: This roadmap contains 75 leetcode questions which will familiarize you with common coding patterns useful in coding interviews

  4. Solving company tagged questions: On leetcode.com you can filter for the company you are interviewing and see top tagged questions for the same. I would recommend solving the top 100 tagged questions based on frequency which are asked in last six months.

  5. Pay attention to design questions: While I recommend focusing on the top 100 questions tagged by the company of interest, interviews often place special emphasis on design-related questions. In these cases, you may be asked to design a class to solve a specific use case. I will go beyond the top 100 questions to find every relevant design question asked in the past year. Here are some top design questions tagged for Meta on Leetcode at the time of writing -

  1. Practice: Once you are done with the above then you can practice timed assessment on leetcode.com for your specific employer or generic ones if not listed. If you want more realistic practice then you can buy some mock interviews on interviewing.io where an engineer from top tech company will take your mock and provide feedback on your performance.
TipUseful tips for coding interviews preparation
  1. The more you practice in conditions like the actual interview—such as using a text editor and working within timed constraints—the better you will perform
  2. When practicing, try to solve the problem on your own for 20-30 minutes before consulting the solution.
  3. If you find a problem challenging to understand, search for [LeetCode Problem #XYZ] on YouTube; you’ll likely find a video with a clearer explanation.
  4. Keeping an Excel sheet to track the problems you’ve solved during practice, along with notes such as ’needs revision, time/space complexity and a summary of the solution, can be very helpful for reviewing later. Refer to Coding Tracking Sheet tab for template in MLE Interview Prep Template.

3 Machine learning System Design (MLSD)

The MLSD interview typically lasts 45 minutes and assesses your ability to solve an abstract ML problem from start to finish. Depending on the level you’re aiming for, you may encounter 1-2 rounds of these interviews. Excalidraw is what is typically used as the platform, but you are just supposed to write / talk about the problem instead of drawing block diagrams in this one. To get a more realistic idea of how a coding interview environment looks like watch this mock interview by interviewing.io -

Harmful Content Removal: Machine Learning (System Design)

3.1 A structured approach to solving machine learning system design problems in interview

In an MLSD interview, it’s crucial for the interviewee to take the lead in the discussion and ensure all aspects of the ML design are covered. The interviews are fast-paced and following a structured method can be highly beneficial:

  1. Clarifying Requirements: In every MLSD interview, you’re typically presented with an abstract problem. For example, you might be asked to “Design a ‘People You May Follow’ recommendation system for Threads or Twitter.” It’s essential to clarify the scope of the problem, ensuring it can be managed within the 45-minute timeframe. Asking the right clarifying questions not only helps you gain clarity but also shows your product awareness. For instance, in the case of a “People You May Follow” recommendation system, you might ask:

    • “Can I assume the purpose of this feature is to help users find influencers or people aligned with their interests?”
    • “On Threads/Twitter, following is unidirectional—one user can follow another without reciprocation. Is that correct?”
    • “What is the estimated total number of users on the platform?
    • “What is the current count of daily active users (DAUs)?”
    • “What’s the average number of people each user follows?”
    • “Since this is a mature platform, can we assume that the user-follow graph is relatively stable and doesn’t change drastically over short periods?”

    Once you are done with the clarifying questions, it’s important to summarize them back to the interviewer, for example: “Okay, we are designing a ‘People You May Follow’ recommendation system for the Threads/Twitter platform. We have around XYZ million daily active users, and we’re assuming that the followership graph remains relatively stable over time.” This helps ensure alignment before diving into the design process.

  2. Frame the problem as an ML task: Once you’ve clarified the requirements, the next step is to map the problem to a known ML objective. This could be something like binary classification, learning to rank, edge prediction, or even visual object detection. Often, a problem can be framed using multiple objectives, and it’s important to explain what each objective aims to achieve and the pros and cons of each approach. For example, in the case of a recommendation system, you could frame it as a learning-to-rank problem or as edge prediction in a user graph. The key is to pick the objective that best aligns with the problem at hand.

  3. Data sources for training labels: Next, brainstorm potential data sources for defining training labels. There could be various options, each with its own trade-offs. For instance, if you’re working on a video recommendation system, you might need to decide between implicit and explicit feedback. Explicit feedback (such as likes, shares, or subscriptions) tends to be high quality but sparse, as not all users provide these signals. On the other hand, implicit feedback (like user’s dwell time or watch time on a post, time spent in the app etc.) is lower quality but available for every user interaction. It’s essential to highlight these trade-offs and choose the best data source for your model’s needs.

  4. Data preparation and Feature Engineering: Once you have identified some potential data sources, it’s good to talk about what features you will create from this data to feed into your ML model. The typical answer revolves around these three “focus area”-

    • Listing entities which are important for the problem: Entities are concepts around which you want to create model features. Some of the typical entities in Meta are users, authors/creators, pages, posts, images, videos, text, events, connections, advertisements, marketplace listings, groups and most importantly user engagement (with these entities).
    • Features we can derive around this entity from the available data sources: For example, for a “User entity” we can utilize many features readily available in user profile table such as their ID, name, age, gender, city, country, language, time zone etc.
    • How to process this data to make it a highly robust and informational feature: For example, user id is a high cardinality feature, and we can convert it into a embedding which our model can learn while user’s gender is a low cardinality feature and can be processed using one-hot encoding.

  5. Model selection: After finishing the data pre-processing step we can talk about choosing the best ML algorithm and architecture for a predictive modeling problem. A typical approach is to go from simple to more complex solutions while explaining its pros and cons. For example, consider a video recommendation system, we can start by establishing a simple baseline by recommending everyone the most popular content, but we will run into situations like no personalization and a bias towards over-represented groups. You can progressively make it complex by talking about collaborative filtering, content-based filtering or a two-tower deep learning model. It’s important to talk about each of these modelling approaches by briefly explaining the algorithm and discussing the tradeoffs. For example, logistic learning might be a good option for you learning a linear task but if the task is complex, we may need to choose more complex models such as ensemble learners or deep learning models.

  6. Model training: Once you have selected the model and defined features, it’s important to talk about some of the details you will consider while training this model. Here are the things to talk about in this section if pertinent to your problem –

    • Splitting data b/w training, testing and evaluation
    • Identifying class imbalanced datasets and techniques to rectify
    • Choosing the right loss function for the model
    • Handling Overfitting and Underfitting

  7. Evaluation: Once you have trained the model, it is important to talk about how you will evaluate the model’s performance. This part can be split into two categories:

    • Offline Evaluation: This is the process of evaluating the model during development phase. An important part here is choosing the right offline metrics to measure how close the predictions are to the ground truth. It’s important to talk about multiple metrics and their pros and cons for the problem in hand. For example, in case of video recommendations there are different metrics to consider like Precision@K, Recall@K, Mean Average Precision(mAP), Normalized Discounted Cumulative Gain (nDCG), Mean Reciprocal Rank (MRR) etc.
    • Online Evaluations: This process involves evaluating the model’s performance in production after deployment. Online metrics, closely aligned with business objectives, typically include multiple measures. As with offline metrics, it is essential to select the appropriate evaluation criteria and provide clear justification for each choice. For instance, in the context of video recommendations, relevant metrics might include the model’s click-through rate (CTR), total watch time, the ratio of completed videos to total views, and the skip rate.

  8. Deployment: Next, if time permits talk about deployment and testing in production. Few things to consider in this section:

    • Does the model need to be deployed online or would batch prediction suffice?
    • How are we going to roll out this change? Running an experiment, or with a phased rollout?
    • How frequently does the model need to be re-trained?
TipTips for MLSD interview:

  • Time Management is Key: Aim to allocate time thoughtfully across all sections of the interview. Keep a mental clock to ensure you provide a well-rounded answer that covers all aspects of ML system design. If you find yourself spending too much time on one section, communicate with the interviewer by saying, “We can revisit this section later if time allows, but I’d like to move on to cover other areas as well.” This demonstrates both maturity and effective time management—qualities crucial at higher levels.

  • Prioritize Practicality Over Perfection: MLSD interviews don’t expect you to design a state-of-the-art (SOTA) solution within 45 minutes. Instead, focus on creating a feasible, well-reasoned solution. Emphasize why you’re choosing specific components and how they contribute to the overall system rather than chasing an ideal but overly complex design.

  • Adapt to Your Interviewer’s Cues: Pay close attention to suggestions from your interviewer. If they ask to deep-dive into a particular section, focus on providing detailed insights into that area. Conversely, if they signal it is time to move on, gracefully wrap up and proceed to the next section. Flexibility and responsiveness highlight your ability to collaborate effectively.

  • Regularly engage with your interviewer: Approach the interviewer as a collaborator in solving the problem. After completing each section, ask for feedback and check whether they would like you to elaborate further or proceed to the next part. This fosters dynamic conversation and ensures you align with their expectations.

3.2 How to prepare for ML system design interview

Preparing for Machine Learning System Design (MLSD) interviews requires a structured approach that combines technical knowledge, system design principles, and practical problem-solving skills. Here’s a step-by-step guide to help you get ready:

  1. Purchase ByteByteGo MLSD Interview book: This book is the only paid resource you need for preparing for MLSD interview. You can buy it from Amazon or you can buy an annual pass on bytebytego website. The first chapter introduces you with ML fundamentals and rest other chapters are focused on ten real ML system design interview questions with detailed solutions. My recommendation is to go thoroughly with the book and read every supplemental links in the reference after every chapter. Once you are done with book, you can refer to these MLSD notes I have created, refer to ML System Design tab in MLE Interview Prep Template.

  2. ML interview prep guide: Meta recruiters provide an interview prep guide which talks about the whole interview experience and MLSD round as well. Its important to read through it and watch the Meta field guide to machine learning series to get more perspective on companies approach to ML.

  3. Meta engineering blog: Reading engineering blog can offer you insights on how MLE’s at Meta have designed some of the key components of most notable features in Meta applications and some of the recent frameworks used internally. Here are some of the most notable blogs are –

  4. Sweat more in training, bleed less in battle: This analogy perfectly captures the essence of preparation for MLSD interviews. The more effort you put into practicing beforehand, the less you’ll need to improvise during the interview. Here are some practical ways to prepare effectively:

    • Design Random Features: Open platforms like Instagram, Facebook, or Threads, pick a random feature, and outline how you would design it from an ML Engineer’s perspective. Consider how the design might differ depending on the app’s unique goals and user base. Here are some examples of features to practice:
      1. Friend suggestions (e.g., “People You May Know”).
      2. News Feed ranking.
      3. Group recommendations.
      4. Short-form video recommendations.
      5. Explore tab content discovery.
      6. Integrity filtering for harmful content.
    • Prepare Features in Advance: As discussed earlier, feature engineering is a crucial part of MLSD interviews. Create a comprehensive list of:
      1. Data sources: Identify potential sources of data (e.g., user activity, content metadata, social graphs).
      2. Entities: Determine key entities like users, items, or sessions.
      3. Features: Brainstorm features around these entities, such as engagement patterns, temporal trends, or semantic embeddings.
      4. Processing Steps: Document how you would pre-process and transform these features for your ML models.

    Preparing these ahead of time will save valuable mental bandwidth during the interview and elevate the quality of your responses. Pre-built templates and examples can help you focus on tailoring solutions to specific problems rather than scrambling to generate ideas under pressure.

  5. Simulate real interview scenarios : Practice solving MLSD problems under realistic conditions by setting a strict time limit (35–40 minutes) and using tools like Excalidraw to create diagrams and workflows. Work through a complete problem, starting from clarifying requirements to designing the entire system. Prioritize clear structure and ensure you address all critical aspects, such as data flow, model selection, architecture, and trade-offs. Recording your sessions can help you pinpoint areas for improvement, such as pacing, clarity, or the depth of your explanations, enabling you to refine your approach and being comfortable with the tools for actual interviews.

4 Behavioral Interviews

The behavioral interview at Meta usually lasts about 45 minutes and focuses on evaluating how well you align with the company’s core values, such as teamwork, leadership, problem-solving, and adaptability. To prepare effectively, it’s important to understand Meta’s culture and values, as well as how to clearly communicate your past experiences. This interview is particularly crucial at higher levels, as it plays a significant role in deciding your level for the position.

4.1 A Structured Approach to Answering Behavioral Questions

One of the best ways to frame your answers is by using the STAR method, which helps you organize your responses clearly and concisely: - Situation: Set the context for your story. Describe the challenge or situation you were in, including the key players and the environment. - Task: Outline your specific responsibility or what you were tasked to do in that situation. - Action: Explain the steps you took to address the task or solve the problem. Focus on the actions you personally took, rather than what the team did as a whole. - Result: Share the outcome of your actions. Quantify the results if possible (e.g., increased efficiency by 20%, reduced costs by 15%) and mention what you learned from the experience.

Using this structure helps you keep your answers clear, logical, and impactful. For example, when asked about teamwork, describe a scenario where you worked with a cross-functional team, the challenge you faced, the actions you took to collaborate, and the successful outcome.

4.2 How to Prepare for Behavioral Interviews

Here’s a step-by-step guide to help you get prepared for Behavioral interviews -

  1. Understand Meta’s Values and Culture: Familiarize yourself with the company’s mission, principle and core values to ensure that your answers reflect how you embody these traits. Think about how your experiences align with Meta’s core values and communicate that through your stories.

  2. Interview Jedi Youtube Video series(Part1, Part2): This series contains top questions asked in behavioral rounds at Meta and covers evaluation criteria, response framework, tips and tricks to prepare + answer them effectively. Having answers pre-prepared to these questions will help you setup for a successful behavioral interview. You can use the Behavioral Tab in MLE Interview Prep Template to write your answers to these questions.

  3. Practice and Refine Your Stories: Once you’ve gone through the videos and written your stories, practice delivering them succinctly using the STAR method. Try to frame each story in a way that highlights your skills and accomplishments while demonstrating your alignment with Meta’s values. Mock interviews with peers or mentors can help you fine-tune your answers and gain confidence in presenting your experiences.

  4. Striking the right balance b/w abstraction and technicality: It’s essential to strike the right balance between abstraction and technicality in your responses. You want to prove your problem-solving and leadership abilities without getting bogged down in excessive technical details. While it’s important to avoid unnecessary technical jargon, you should still provide sufficient detail to show the depth of your involvement. For example, instead of providing internal abbreviations like ECR say managed Docker container registry. Watch this section from A Life Engineered Video to understand the difference b/w answer to the same behavioral question at different levels.

  5. Be Honest and Authentic: Meta values authenticity and integrity, so avoid overselling yourself. Be honest about your experiences, especially when discussing challenges or failures. What matters most is your ability to reflect on your experiences, learn from them, and demonstrate how you’ve grown.

By following these strategies, you can approach Meta’s behavioral interviews with a structured, thoughtful approach that showcases your qualifications and demonstrates your alignment with the company’s values.

5 Closing Thoughts

Preparing for the Meta MLE interview can feel like a challenging journey, but with the right preparation and resources, you can navigate it successfully. Lastly, remember that the process itself is a learning experience. Each interview is an opportunity to grow, learn from feedback, and refine your skills. Whether you land the job or not, the preparation you put in will enhance your overall growth as a machine learning engineer.

Stay focused, stay curious, and approach each step with confidence. Good luck!

I hope you enjoyed reading it. If there is any feedback on the code or just the blog post, feel free to comment below or reach out on LinkedIn.

]]> ML Interview Guide https://aayushmnit.com/posts/2024-12-15-MLInterviewPrep/MLInterviewPrep.html Sun, 15 Dec 2024 08:00:00 GMT privateGPT Walkthrough Aayush Agrawal https://aayushmnit.com/posts/2023-05-22-PrivateGPTWalkthrough/privateGPTWalkthrough.html

A code walkthrough of privateGPT repo on how to build your own offline GPT Q&A system.

Large Language Models (LLMs) have surged in popularity, pushing the boundaries of natural language processing. OpenAI’s GPT-3.5 is a prime example, revolutionizing our technology interactions and sparking innovation. Particularly, LLMs excel in building Question Answering applications on knowledge bases. In this blog, we delve into the top trending GitHub repository for this week: the PrivateGPT repository and do a code walkthrough.

Fig. 1: Private GPT on GitHub’s top trending chart

1 What is privateGPT?

One of the primary concerns associated with employing online interfaces like OpenAI chatGPT or other Large Language Model systems pertains to data privacy, data control, and potential data leakage. The privateGPT repository presents a fully offline alternative for engaging with personal documents. It is constructed using open source tools and technology, thereby enabling the utilization of LLMs capabilities without compromising data privacy or encountering data leakage issues.

Fig.2: privateGPT on GitHub. At the time of writing repo had 19K+ stars and 2k+ forks.

2 Running privateGPT locally

To run privateGPT locally, users need to install the necessary packages, configure specific variables, and provide their knowledge base for question-answering purposes. Additional information on the installation process and usage can be found in the repository documentation or by referring to a dedicated blog post on the topic.

Essentially you can run it by calling the privateGPT.py file like -

python privateGPT.py
Fig.3: Invoking privateGPT locally and asking a question.

And get a response that also mention the sources it looked up for context.

Fig.4: privateGPT response.

3 Code Walkthrough

privateGPT code comprises two pipelines:

  1. Ingestion Pipeline: This pipeline is responsible for converting and storing your documents, as well as generating embeddings for them. The documents are stored in a suitable format, and their embeddings are stored in an embedding database.

  2. Q&A Interface: This interface accepts user prompts, the embedding database, and an open-source Language Model (LM) model as inputs. It utilizes these inputs to generate responses to the user’s queries.

3.1 Ingestion Pipeline

Let’s delve into the ingestion pipeline for a closer examination. The ingestion pipeline encompasses the following steps:

  1. Identifying files with various extensions and retrieving all the knowledge base from the source directory.

  2. Splitting the documents into smaller chunks based on the parameters of chunk_size and chunk_overlap.

  3. Initializing the Huggingfaceembeddings module of langchain. This involves loading a pre-trained language model from the sentence_transformers library.

  4. Initializing the Chroma database from langchain.vectorstores. This step involves taking the chunked text and the initialized embedding model and saving it in the embedding database on disk.

Fig.5: Ingestion Pipeline

Let’s look at these steps one by one.

3.1.1 Identifying and loading files from the source directory

First, we import the required libraries and various text loaders from langchain.document_loaders.

import os
import glob
from typing import List
from multiprocessing import Pool
from tqdm import tqdm
from langchain.document_loaders import (
    CSVLoader,
    EverNoteLoader,
    PDFMinerLoader,
    TextLoader,
    UnstructuredEmailLoader,
    UnstructuredEPubLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredODTLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader,
)
from langchain.docstore.document import Document

Next, we define the mapping b/w each extension and their respective langchain document loader. You can read document loader documentation for more available loaders.

# Map file extensions to document loaders and their arguments
LOADER_MAPPING = {
    ".csv": (CSVLoader, {}),
    ".doc": (UnstructuredWordDocumentLoader, {}),
    ".docx": (UnstructuredWordDocumentLoader, {}),
    ".enex": (EverNoteLoader, {}),
    ".epub": (UnstructuredEPubLoader, {}),
    ".html": (UnstructuredHTMLLoader, {}),
    ".md": (UnstructuredMarkdownLoader, {}),
    ".odt": (UnstructuredODTLoader, {}),
    ".pdf": (PDFMinerLoader, {}),
    ".ppt": (UnstructuredPowerPointLoader, {}),
    ".pptx": (UnstructuredPowerPointLoader, {}),
    ".txt": (TextLoader, {"encoding": "utf8"}),
}

Next, we define our single document loader.

def load_single_document(file_path: str) -> Document:
    ## Find extension of the file
    ext = "." + file_path.rsplit(".", 1)[-1] 
    if ext in LOADER_MAPPING: 
        # Find the appropriate loader class and arguments
        loader_class, loader_args = LOADER_MAPPING[ext] 
        # Invoke the instance of document loader
        loader = loader_class(file_path, **loader_args) 
        ## Return the loaded document
        return loader.load()[0] 
    raise ValueError(f"Unsupported file extension '{ext}'")
    
git_dir = "../../../../privateGPT/"
loaded_document = load_single_document(git_dir+'source_documents/state_of_the_union.txt')
print(f'Type of loaded document {type(loaded_document)}')
loaded_document
Type of loaded document <class 'langchain.schema.Document'>
Document(page_content='Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.  \n\nLast year COVID-19 kept us apart. This year we are finally together again. \n\nTonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \n\nWith a duty to one another to the American people to the Constitution. \n\nAnd with an unwavering resolve that freedom will always triumph over tyranny. \n\nSix days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. \n\nHe thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. \n\nHe met the Ukrainian people. \n\nFrom President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world. \n\nGroups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. \n\nIn this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight. \n\nLet each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world. \n\nPlease rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people. \n\nThroughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos.   \n\nThey keep moving.   \n\nAnd the costs and the threats to America and the world keep rising.   \n\nThat’s why the NATO Alliance was created to secure peace and stability in Europe after World War 2. \n\nThe United States is a member along with 29 other nations. \n\nIt matters. American diplomacy matters. American resolve matters. \n\nPutin’s latest attack on Ukraine was premeditated and unprovoked. \n\nHe rejected repeated efforts at diplomacy. \n\nHe thought the West and NATO wouldn’t respond. And he thought he could divide us at home. Putin was wrong. We were ready.  Here is what we did.   \n\nWe prepared extensively and carefully. \n\nWe spent months building a coalition of other freedom-loving nations from Europe and the Americas to Asia and Africa to confront Putin. \n\nI spent countless hours unifying our European allies. We shared with the world in advance what we knew Putin was planning and precisely how he would try to falsely justify his aggression.  \n\nWe countered Russia’s lies with truth.   \n\nAnd now that he has acted the free world is holding him accountable. \n\nAlong with twenty-seven members of the European Union including France, Germany, Italy, as well as countries like the United Kingdom, Canada, Japan, Korea, Australia, New Zealand, and many others, even Switzerland. \n\nWe are inflicting pain on Russia and supporting the people of Ukraine. Putin is now isolated from the world more than ever. \n\nTogether with our allies –we are right now enforcing powerful economic sanctions. \n\nWe are cutting off Russia’s largest banks from the international financial system.  \n\nPreventing Russia’s central bank from defending the Russian Ruble making Putin’s $630 Billion “war fund” worthless.   \n\nWe are choking off Russia’s access to technology that will sap its economic strength and weaken its military for years to come.  \n\nTonight I say to the Russian oligarchs and corrupt leaders who have bilked billions of dollars off this violent regime no more. \n\nThe U.S. Department of Justice is assembling a dedicated task force to go after the crimes of Russian oligarchs.  \n\nWe are joining with our European allies to find and seize your yachts your luxury apartments your private jets. We are coming for your ill-begotten gains. \n\nAnd tonight I am announcing that we will join our allies in closing off American air space to all Russian flights – further isolating Russia – and adding an additional squeeze –on their economy. The Ruble has lost 30% of its value. \n\nThe Russian stock market has lost 40% of its value and trading remains suspended. Russia’s economy is reeling and Putin alone is to blame. \n\nTogether with our allies we are providing support to the Ukrainians in their fight for freedom. Military assistance. Economic assistance. Humanitarian assistance. \n\nWe are giving more than $1 Billion in direct assistance to Ukraine. \n\nAnd we will continue to aid the Ukrainian people as they defend their country and to help ease their suffering.  \n\nLet me be clear, our forces are not engaged and will not engage in conflict with Russian forces in Ukraine.  \n\nOur forces are not going to Europe to fight in Ukraine, but to defend our NATO Allies – in the event that Putin decides to keep moving west.  \n\nFor that purpose we’ve mobilized American ground forces, air squadrons, and ship deployments to protect NATO countries including Poland, Romania, Latvia, Lithuania, and Estonia. \n\nAs I have made crystal clear the United States and our Allies will defend every inch of territory of NATO countries with the full force of our collective power.  \n\nAnd we remain clear-eyed. The Ukrainians are fighting back with pure courage. But the next few days weeks, months, will be hard on them.  \n\nPutin has unleashed violence and chaos.  But while he may make gains on the battlefield – he will pay a continuing high price over the long run. \n\nAnd a proud Ukrainian people, who have known 30 years  of independence, have repeatedly shown that they will not tolerate anyone who tries to take their country backwards.  \n\nTo all Americans, I will be honest with you, as I’ve always promised. A Russian dictator, invading a foreign country, has costs around the world. \n\nAnd I’m taking robust action to make sure the pain of our sanctions  is targeted at Russia’s economy. And I will use every tool at our disposal to protect American businesses and consumers. \n\nTonight, I can announce that the United States has worked with 30 other countries to release 60 Million barrels of oil from reserves around the world.  \n\nAmerica will lead that effort, releasing 30 Million barrels from our own Strategic Petroleum Reserve. And we stand ready to do more if necessary, unified with our allies.  \n\nThese steps will help blunt gas prices here at home. And I know the news about what’s happening can seem alarming. \n\nBut I want you to know that we are going to be okay. \n\nWhen the history of this era is written Putin’s war on Ukraine will have left Russia weaker and the rest of the world stronger. \n\nWhile it shouldn’t have taken something so terrible for people around the world to see what’s at stake now everyone sees it clearly. \n\nWe see the unity among leaders of nations and a more unified Europe a more unified West. And we see unity among the people who are gathering in cities in large crowds around the world even in Russia to demonstrate their support for Ukraine.  \n\nIn the battle between democracy and autocracy, democracies are rising to the moment, and the world is clearly choosing the side of peace and security. \n\nThis is a real test. It’s going to take time. So let us continue to draw inspiration from the iron will of the Ukrainian people. \n\nTo our fellow Ukrainian Americans who forge a deep bond that connects our two nations we stand with you. \n\nPutin may circle Kyiv with tanks, but he will never gain the hearts and souls of the Ukrainian people. \n\nHe will never extinguish their love of freedom. He will never weaken the resolve of the free world. \n\nWe meet tonight in an America that has lived through two of the hardest years this nation has ever faced. \n\nThe pandemic has been punishing. \n\nAnd so many families are living paycheck to paycheck, struggling to keep up with the rising cost of food, gas, housing, and so much more. \n\nI understand. \n\nI remember when my Dad had to leave our home in Scranton, Pennsylvania to find work. I grew up in a family where if the price of food went up, you felt it. \n\nThat’s why one of the first things I did as President was fight to pass the American Rescue Plan.  \n\nBecause people were hurting. We needed to act, and we did. \n\nFew pieces of legislation have done more in a critical moment in our history to lift us out of crisis. \n\nIt fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans.  \n\nHelped put food on their table, keep a roof over their heads, and cut the cost of health insurance. \n\nAnd as my Dad used to say, it gave people a little breathing room. \n\nAnd unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind. \n\nAnd it worked. It created jobs. Lots of jobs. \n\nIn fact—our economy created over 6.5 Million new jobs just last year, more jobs created in one year  \nthan ever before in the history of America. \n\nOur economy grew at a rate of 5.7% last year, the strongest growth in nearly 40 years, the first step in bringing fundamental change to an economy that hasn’t worked for the working people of this nation for too long.  \n\nFor the past 40 years we were told that if we gave tax breaks to those at the very top, the benefits would trickle down to everyone else. \n\nBut that trickle-down theory led to weaker economic growth, lower wages, bigger deficits, and the widest gap between those at the top and everyone else in nearly a century. \n\nVice President Harris and I ran for office with a new economic vision for America. \n\nInvest in America. Educate Americans. Grow the workforce. Build the economy from the bottom up  \nand the middle out, not from the top down.  \n\nBecause we know that when the middle class grows, the poor have a ladder up and the wealthy do very well. \n\nAmerica used to have the best roads, bridges, and airports on Earth. \n\nNow our infrastructure is ranked 13th in the world. \n\nWe won’t be able to compete for the jobs of the 21st Century if we don’t fix that. \n\nThat’s why it was so important to pass the Bipartisan Infrastructure Law—the most sweeping investment to rebuild America in history. \n\nThis was a bipartisan effort, and I want to thank the members of both parties who worked to make it happen. \n\nWe’re done talking about infrastructure weeks. \n\nWe’re going to have an infrastructure decade. \n\nIt is going to transform America and put us on a path to win the economic competition of the 21st Century that we face with the rest of the world—particularly with China.  \n\nAs I’ve told Xi Jinping, it is never a good bet to bet against the American people. \n\nWe’ll create good jobs for millions of Americans, modernizing roads, airports, ports, and waterways all across America. \n\nAnd we’ll do it all to withstand the devastating effects of the climate crisis and promote environmental justice. \n\nWe’ll build a national network of 500,000 electric vehicle charging stations, begin to replace poisonous lead pipes—so every child—and every American—has clean water to drink at home and at school, provide affordable high-speed internet for every American—urban, suburban, rural, and tribal communities. \n\n4,000 projects have already been announced. \n\nAnd tonight, I’m announcing that this year we will start fixing over 65,000 miles of highway and 1,500 bridges in disrepair. \n\nWhen we use taxpayer dollars to rebuild America – we are going to Buy American: buy American products to support American jobs. \n\nThe federal government spends about $600 Billion a year to keep the country safe and secure. \n\nThere’s been a law on the books for almost a century \nto make sure taxpayers’ dollars support American jobs and businesses. \n\nEvery Administration says they’ll do it, but we are actually doing it. \n\nWe will buy American to make sure everything from the deck of an aircraft carrier to the steel on highway guardrails are made in America. \n\nBut to compete for the best jobs of the future, we also need to level the playing field with China and other competitors. \n\nThat’s why it is so important to pass the Bipartisan Innovation Act sitting in Congress that will make record investments in emerging technologies and American manufacturing. \n\nLet me give you one example of why it’s so important to pass it. \n\nIf you travel 20 miles east of Columbus, Ohio, you’ll find 1,000 empty acres of land. \n\nIt won’t look like much, but if you stop and look closely, you’ll see a “Field of dreams,” the ground on which America’s future will be built. \n\nThis is where Intel, the American company that helped build Silicon Valley, is going to build its $20 billion semiconductor “mega site”. \n\nUp to eight state-of-the-art factories in one place. 10,000 new good-paying jobs. \n\nSome of the most sophisticated manufacturing in the world to make computer chips the size of a fingertip that power the world and our everyday lives. \n\nSmartphones. The Internet. Technology we have yet to invent. \n\nBut that’s just the beginning. \n\nIntel’s CEO, Pat Gelsinger, who is here tonight, told me they are ready to increase their investment from  \n$20 billion to $100 billion. \n\nThat would be one of the biggest investments in manufacturing in American history. \n\nAnd all they’re waiting for is for you to pass this bill. \n\nSo let’s not wait any longer. Send it to my desk. I’ll sign it.  \n\nAnd we will really take off. \n\nAnd Intel is not alone. \n\nThere’s something happening in America. \n\nJust look around and you’ll see an amazing story. \n\nThe rebirth of the pride that comes from stamping products “Made In America.” The revitalization of American manufacturing.   \n\nCompanies are choosing to build new factories here, when just a few years ago, they would have built them overseas. \n\nThat’s what is happening. Ford is investing $11 billion to build electric vehicles, creating 11,000 jobs across the country. \n\nGM is making the largest investment in its history—$7 billion to build electric vehicles, creating 4,000 jobs in Michigan. \n\nAll told, we created 369,000 new manufacturing jobs in America just last year. \n\nPowered by people I’ve met like JoJo Burgess, from generations of union steelworkers from Pittsburgh, who’s here with us tonight. \n\nAs Ohio Senator Sherrod Brown says, “It’s time to bury the label “Rust Belt.” \n\nIt’s time. \n\nBut with all the bright spots in our economy, record job growth and higher wages, too many families are struggling to keep up with the bills.  \n\nInflation is robbing them of the gains they might otherwise feel. \n\nI get it. That’s why my top priority is getting prices under control. \n\nLook, our economy roared back faster than most predicted, but the pandemic meant that businesses had a hard time hiring enough workers to keep up production in their factories. \n\nThe pandemic also disrupted global supply chains. \n\nWhen factories close, it takes longer to make goods and get them from the warehouse to the store, and prices go up. \n\nLook at cars. \n\nLast year, there weren’t enough semiconductors to make all the cars that people wanted to buy. \n\nAnd guess what, prices of automobiles went up. \n\nSo—we have a choice. \n\nOne way to fight inflation is to drive down wages and make Americans poorer.  \n\nI have a better plan to fight inflation. \n\nLower your costs, not your wages. \n\nMake more cars and semiconductors in America. \n\nMore infrastructure and innovation in America. \n\nMore goods moving faster and cheaper in America. \n\nMore jobs where you can earn a good living in America. \n\nAnd instead of relying on foreign supply chains, let’s make it in America. \n\nEconomists call it “increasing the productive capacity of our economy.” \n\nI call it building a better America. \n\nMy plan to fight inflation will lower your costs and lower the deficit. \n\n17 Nobel laureates in economics say my plan will ease long-term inflationary pressures. Top business leaders and most Americans support my plan. And here’s the plan: \n\nFirst – cut the cost of prescription drugs. Just look at insulin. One in ten Americans has diabetes. In Virginia, I met a 13-year-old boy named Joshua Davis.  \n\nHe and his Dad both have Type 1 diabetes, which means they need insulin every day. Insulin costs about $10 a vial to make.  \n\nBut drug companies charge families like Joshua and his Dad up to 30 times more. I spoke with Joshua’s mom. \n\nImagine what it’s like to look at your child who needs insulin and have no idea how you’re going to pay for it.  \n\nWhat it does to your dignity, your ability to look your child in the eye, to be the parent you expect to be. \n\nJoshua is here with us tonight. Yesterday was his birthday. Happy birthday, buddy.  \n\nFor Joshua, and for the 200,000 other young people with Type 1 diabetes, let’s cap the cost of insulin at $35 a month so everyone can afford it.  \n\nDrug companies will still do very well. And while we’re at it let Medicare negotiate lower prices for prescription drugs, like the VA already does. \n\nLook, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. \n\nSecond – cut energy costs for families an average of $500 a year by combatting climate change.  \n\nLet’s provide investments and tax credits to weatherize your homes and businesses to be energy efficient and you get a tax credit; double America’s clean energy production in solar, wind, and so much more;  lower the price of electric vehicles, saving you another $80 a month because you’ll never have to pay at the gas pump again. \n\nThird – cut the cost of child care. Many families pay up to $14,000 a year for child care per child.  \n\nMiddle-class and working families shouldn’t have to pay more than 7% of their income for care of young children.  \n\nMy plan will cut the cost in half for most families and help parents, including millions of women, who left the workforce during the pandemic because they couldn’t afford child care, to be able to get back to work. \n\nMy plan doesn’t stop there. It also includes home and long-term care. More affordable housing. And Pre-K for every 3- and 4-year-old.  \n\nAll of these will lower costs. \n\nAnd under my plan, nobody earning less than $400,000 a year will pay an additional penny in new taxes. Nobody.  \n\nThe one thing all Americans agree on is that the tax system is not fair. We have to fix it.  \n\nI’m not looking to punish anyone. But let’s make sure corporations and the wealthiest Americans start paying their fair share. \n\nJust last year, 55 Fortune 500 corporations earned $40 billion in profits and paid zero dollars in federal income tax.  \n\nThat’s simply not fair. That’s why I’ve proposed a 15% minimum tax rate for corporations. \n\nWe got more than 130 countries to agree on a global minimum tax rate so companies can’t get out of paying their taxes at home by shipping jobs and factories overseas. \n\nThat’s why I’ve proposed closing loopholes so the very wealthy don’t pay a lower tax rate than a teacher or a firefighter.  \n\nSo that’s my plan. It will grow the economy and lower costs for families. \n\nSo what are we waiting for? Let’s get this done. And while you’re at it, confirm my nominees to the Federal Reserve, which plays a critical role in fighting inflation.  \n\nMy plan will not only lower costs to give families a fair shot, it will lower the deficit. \n\nThe previous Administration not only ballooned the deficit with tax cuts for the very wealthy and corporations, it undermined the watchdogs whose job was to keep pandemic relief funds from being wasted. \n\nBut in my administration, the watchdogs have been welcomed back. \n\nWe’re going after the criminals who stole billions in relief money meant for small businesses and millions of Americans.  \n\nAnd tonight, I’m announcing that the Justice Department will name a chief prosecutor for pandemic fraud. \n\nBy the end of this year, the deficit will be down to less than half what it was before I took office.  \n\nThe only president ever to cut the deficit by more than one trillion dollars in a single year. \n\nLowering your costs also means demanding more competition. \n\nI’m a capitalist, but capitalism without competition isn’t capitalism. \n\nIt’s exploitation—and it drives up prices. \n\nWhen corporations don’t have to compete, their profits go up, your prices go up, and small businesses and family farmers and ranchers go under. \n\nWe see it happening with ocean carriers moving goods in and out of America. \n\nDuring the pandemic, these foreign-owned companies raised prices by as much as 1,000% and made record profits. \n\nTonight, I’m announcing a crackdown on these companies overcharging American businesses and consumers. \n\nAnd as Wall Street firms take over more nursing homes, quality in those homes has gone down and costs have gone up.  \n\nThat ends on my watch. \n\nMedicare is going to set higher standards for nursing homes and make sure your loved ones get the care they deserve and expect. \n\nWe’ll also cut costs and keep the economy going strong by giving workers a fair shot, provide more training and apprenticeships, hire them based on their skills not degrees. \n\nLet’s pass the Paycheck Fairness Act and paid leave.  \n\nRaise the minimum wage to $15 an hour and extend the Child Tax Credit, so no one has to raise a family in poverty. \n\nLet’s increase Pell Grants and increase our historic support of HBCUs, and invest in what Jill—our First Lady who teaches full-time—calls America’s best-kept secret: community colleges. \n\nAnd let’s pass the PRO Act when a majority of workers want to form a union—they shouldn’t be stopped.  \n\nWhen we invest in our workers, when we build the economy from the bottom up and the middle out together, we can do something we haven’t done in a long time: build a better America. \n\nFor more than two years, COVID-19 has impacted every decision in our lives and the life of the nation. \n\nAnd I know you’re tired, frustrated, and exhausted. \n\nBut I also know this. \n\nBecause of the progress we’ve made, because of your resilience and the tools we have, tonight I can say  \nwe are moving forward safely, back to more normal routines.  \n\nWe’ve reached a new moment in the fight against COVID-19, with severe cases down to a level not seen since last July.  \n\nJust a few days ago, the Centers for Disease Control and Prevention—the CDC—issued new mask guidelines. \n\nUnder these new guidelines, most Americans in most of the country can now be mask free.   \n\nAnd based on the projections, more of the country will reach that point across the next couple of weeks. \n\nThanks to the progress we have made this past year, COVID-19 need no longer control our lives.  \n\nI know some are talking about “living with COVID-19”. Tonight – I say that we will never just accept living with COVID-19. \n\nWe will continue to combat the virus as we do other diseases. And because this is a virus that mutates and spreads, we will stay on guard. \n\nHere are four common sense steps as we move forward safely.  \n\nFirst, stay protected with vaccines and treatments. We know how incredibly effective vaccines are. If you’re vaccinated and boosted you have the highest degree of protection. \n\nWe will never give up on vaccinating more Americans. Now, I know parents with kids under 5 are eager to see a vaccine authorized for their children. \n\nThe scientists are working hard to get that done and we’ll be ready with plenty of vaccines when they do. \n\nWe’re also ready with anti-viral treatments. If you get COVID-19, the Pfizer pill reduces your chances of ending up in the hospital by 90%.  \n\nWe’ve ordered more of these pills than anyone in the world. And Pfizer is working overtime to get us 1 Million pills this month and more than double that next month.  \n\nAnd we’re launching the “Test to Treat” initiative so people can get tested at a pharmacy, and if they’re positive, receive antiviral pills on the spot at no cost.  \n\nIf you’re immunocompromised or have some other vulnerability, we have treatments and free high-quality masks. \n\nWe’re leaving no one behind or ignoring anyone’s needs as we move forward. \n\nAnd on testing, we have made hundreds of millions of tests available for you to order for free.   \n\nEven if you already ordered free tests tonight, I am announcing that you can order more from covidtests.gov starting next week. \n\nSecond – we must prepare for new variants. Over the past year, we’ve gotten much better at detecting new variants. \n\nIf necessary, we’ll be able to deploy new vaccines within 100 days instead of many more months or years.  \n\nAnd, if Congress provides the funds we need, we’ll have new stockpiles of tests, masks, and pills ready if needed. \n\nI cannot promise a new variant won’t come. But I can promise you we’ll do everything within our power to be ready if it does.  \n\nThird – we can end the shutdown of schools and businesses. We have the tools we need. \n\nIt’s time for Americans to get back to work and fill our great downtowns again.  People working from home can feel safe to begin to return to the office.   \n\nWe’re doing that here in the federal government. The vast majority of federal workers will once again work in person. \n\nOur schools are open. Let’s keep it that way. Our kids need to be in school. \n\nAnd with 75% of adult Americans fully vaccinated and hospitalizations down by 77%, most Americans can remove their masks, return to work, stay in the classroom, and move forward safely. \n\nWe achieved this because we provided free vaccines, treatments, tests, and masks. \n\nOf course, continuing this costs money. \n\nI will soon send Congress a request. \n\nThe vast majority of Americans have used these tools and may want to again, so I expect Congress to pass it quickly.   \n\nFourth, we will continue vaccinating the world.     \n\nWe’ve sent 475 Million vaccine doses to 112 countries, more than any other nation. \n\nAnd we won’t stop. \n\nWe have lost so much to COVID-19. Time with one another. And worst of all, so much loss of life. \n\nLet’s use this moment to reset. Let’s stop looking at COVID-19 as a partisan dividing line and see it for what it is: A God-awful disease.  \n\nLet’s stop seeing each other as enemies, and start seeing each other for who we really are: Fellow Americans.  \n\nWe can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. \n\nI recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. \n\nThey were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. \n\nOfficer Mora was 27 years old. \n\nOfficer Rivera was 22. \n\nBoth Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. \n\nI spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \n\nI’ve worked on these issues a long time. \n\nI know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. \n\nSo let’s not abandon our streets. Or choose between safety and equal justice. \n\nLet’s come together to protect our communities, restore trust, and hold law enforcement accountable. \n\nThat’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. \n\nThat’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.  \n\nWe should all agree: The answer is not to Defund the police. The answer is to FUND the police with the resources and training they need to protect our communities. \n\nI ask Democrats and Republicans alike: Pass my budget and keep our neighborhoods safe.  \n\nAnd I will keep doing everything in my power to crack down on gun trafficking and ghost guns you can buy online and make at home—they have no serial numbers and can’t be traced. \n\nAnd I ask Congress to pass proven measures to reduce gun violence. Pass universal background checks. Why should anyone on a terrorist list be able to purchase a weapon? \n\nBan assault weapons and high-capacity magazines. \n\nRepeal the liability shield that makes gun manufacturers the only industry in America that can’t be sued. \n\nThese laws don’t infringe on the Second Amendment. They save lives. \n\nThe most fundamental right in America is the right to vote – and to have it counted. And it’s under assault. \n\nIn state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. \n\nWe cannot let this happen. \n\nTonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n\nTonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n\nOne of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n\nAnd I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. \n\nA former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling.  \n\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers.  \n\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders. \n\nWe can do all this while keeping lit the torch of liberty that has led generations of immigrants to this land—my forefathers and so many of yours. \n\nProvide a pathway to citizenship for Dreamers, those on temporary status, farm workers, and essential workers. \n\nRevise our laws so businesses have the workers they need and families don’t wait decades to reunite. \n\nIt’s not only the right thing to do—it’s the economically smart thing to do. \n\nThat’s why immigration reform is supported by everyone from labor unions to religious leaders to the U.S. Chamber of Commerce. \n\nLet’s get it done once and for all. \n\nAdvancing liberty and justice also requires protecting the rights of women. \n\nThe constitutional right affirmed in Roe v. Wade—standing precedent for half a century—is under attack as never before. \n\nIf we want to go forward—not backward—we must protect access to health care. Preserve a woman’s right to choose. And let’s continue to advance maternal health care in America. \n\nAnd for our LGBTQ+ Americans, let’s finally get the bipartisan Equality Act to my desk. The onslaught of state laws targeting transgender Americans and their families is wrong. \n\nAs I said last year, especially to our younger transgender Americans, I will always have your back as your President, so you can be yourself and reach your God-given potential. \n\nWhile it often appears that we never agree, that isn’t true. I signed 80 bipartisan bills into law last year. From preventing government shutdowns to protecting Asian-Americans from still-too-common hate crimes to reforming military justice. \n\nAnd soon, we’ll strengthen the Violence Against Women Act that I first wrote three decades ago. It is important for us to show the nation that we can come together and do big things. \n\nSo tonight I’m offering a Unity Agenda for the Nation. Four big things we can do together.  \n\nFirst, beat the opioid epidemic. \n\nThere is so much we can do. Increase funding for prevention, treatment, harm reduction, and recovery.  \n\nGet rid of outdated rules that stop doctors from prescribing treatments. And stop the flow of illicit drugs by working with state and local law enforcement to go after traffickers. \n\nIf you’re suffering from addiction, know you are not alone. I believe in recovery, and I celebrate the 23 million Americans in recovery. \n\nSecond, let’s take on mental health. Especially among our children, whose lives and education have been turned upside down.  \n\nThe American Rescue Plan gave schools money to hire teachers and help students make up for lost learning.  \n\nI urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. \n\nChildren were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media. \n\nAs Frances Haugen, who is here with us tonight, has shown, we must hold social media platforms accountable for the national experiment they’re conducting on our children for profit. \n\nIt’s time to strengthen privacy protections, ban targeted advertising to children, demand tech companies stop collecting personal data on our children. \n\nAnd let’s get all Americans the mental health services they need. More people they can turn to for help, and full parity between physical and mental health care. \n\nThird, support our veterans. \n\nVeterans are the best of us. \n\nI’ve always believed that we have a sacred obligation to equip all those we send to war and care for them and their families when they come home. \n\nMy administration is providing assistance with job training and housing, and now helping lower-income veterans get VA care debt-free.  \n\nOur troops in Iraq and Afghanistan faced many dangers. \n\nOne was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. \n\nWhen they came home, many of the world’s fittest and best trained warriors were never the same. \n\nHeadaches. Numbness. Dizziness. \n\nA cancer that would put them in a flag-draped coffin. \n\nI know. \n\nOne of those soldiers was my son Major Beau Biden. \n\nWe don’t know for sure if a burn pit was the cause of his brain cancer, or the diseases of so many of our troops. \n\nBut I’m committed to finding out everything we can. \n\nCommitted to military families like Danielle Robinson from Ohio. \n\nThe widow of Sergeant First Class Heath Robinson.  \n\nHe was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. \n\nStationed near Baghdad, just yards from burn pits the size of football fields. \n\nHeath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter. \n\nBut cancer from prolonged exposure to burn pits ravaged Heath’s lungs and body. \n\nDanielle says Heath was a fighter to the very end. \n\nHe didn’t know how to stop fighting, and neither did she. \n\nThrough her pain she found purpose to demand we do better. \n\nTonight, Danielle—we are. \n\nThe VA is pioneering new ways of linking toxic exposures to diseases, already helping more veterans get benefits. \n\nAnd tonight, I’m announcing we’re expanding eligibility to veterans suffering from nine respiratory cancers. \n\nI’m also calling on Congress: pass a law to make sure veterans devastated by toxic exposures in Iraq and Afghanistan finally get the benefits and comprehensive health care they deserve. \n\nAnd fourth, let’s end cancer as we know it. \n\nThis is personal to me and Jill, to Kamala, and to so many of you. \n\nCancer is the #2 cause of death in America–second only to heart disease. \n\nLast month, I announced our plan to supercharge  \nthe Cancer Moonshot that President Obama asked me to lead six years ago. \n\nOur goal is to cut the cancer death rate by at least 50% over the next 25 years, turn more cancers from death sentences into treatable diseases.  \n\nMore support for patients and families. \n\nTo get there, I call on Congress to fund ARPA-H, the Advanced Research Projects Agency for Health. \n\nIt’s based on DARPA—the Defense Department project that led to the Internet, GPS, and so much more.  \n\nARPA-H will have a singular purpose—to drive breakthroughs in cancer, Alzheimer’s, diabetes, and more. \n\nA unity agenda for the nation. \n\nWe can do this. \n\nMy fellow Americans—tonight , we have gathered in a sacred space—the citadel of our democracy. \n\nIn this Capitol, generation after generation, Americans have debated great questions amid great strife, and have done great things. \n\nWe have fought for freedom, expanded liberty, defeated totalitarianism and terror. \n\nAnd built the strongest, freest, and most prosperous nation the world has ever known. \n\nNow is the hour. \n\nOur moment of responsibility. \n\nOur test of resolve and conscience, of history itself. \n\nIt is in this moment that our character is formed. Our purpose is found. Our future is forged. \n\nWell I know this nation.  \n\nWe will meet the test. \n\nTo protect freedom and liberty, to expand fairness and opportunity. \n\nWe will save democracy. \n\nAs hard as these times have been, I am more optimistic about America today than I have been my whole life. \n\nBecause I see the future that is within our grasp. \n\nBecause I know there is simply nothing beyond our capacity. \n\nWe are the only nation on Earth that has always turned every crisis we have faced into an opportunity. \n\nThe only nation that can be defined by a single word: possibilities. \n\nSo on this night, in our 245th year as a nation, I have come to report on the State of the Union. \n\nAnd my report is this: the State of the Union is strong—because you, the American people, are strong. \n\nWe are stronger today than we were a year ago. \n\nAnd we will be stronger a year from now than we are today. \n\nNow is our moment to meet and overcome the challenges of our time. \n\nAnd we will, as one people. \n\nOne America. \n\nThe United States of America. \n\nMay God bless you all. May God protect our troops.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'})

The load_single_document function accomplishes the following steps:

  1. Extracts the file extension from the given file path.
  2. Retrieves the corresponding document loader and its arguments from the previously defined LOADER_MAPPING dictionary.
  3. Creates an instance of the appropriate document loader.
  4. Loads the document using the instantiated loader.
  5. Returns the loaded document.

We can see that load_single_document returns a document of type langchain.schema.Document. Which according to the documentation consists of page_content (the content of the data) and metadata (auxiliary pieces of information describing attributes of the data).

def load_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
    """
    Loads all documents from the source documents directory, ignoring specified files
    """
    all_files = []
    for ext in LOADER_MAPPING:
        #Find all the files within source documents which matches the extensions in Loader_Mapping file
        all_files.extend(
            glob.glob(os.path.join(source_dir, f"**/*{ext}"), recursive=True)
        )
    
    ## Filtering files from all_files if its in ignored_files
    filtered_files = [file_path for file_path in all_files if file_path not in ignored_files]
    
    ## Spinning up resource pool
    with Pool(processes=os.cpu_count()) as pool:
        results = []
        with tqdm(total=len(filtered_files), desc='Loading new documents', ncols=80) as pbar:
            # Load each document from filtered files list using load_single_document function
            for i, doc in enumerate(pool.imap_unordered(load_single_document, filtered_files)):
                results.append(doc)
                pbar.update()
    
    return results

The load_single_documents function carries out the following steps: 1. Initializes an empty dictionary called all_files.
2. For each extension in the LOADER_MAPPING dictionary, it searches for all the files with that extension in the source directory and adds them to the all_files list.
3. Creates a new list named filtered_files by removing the files listed in the ignored_files list from the all_files list.
4. Executes a parallel loading operation on all the files in the filtered_files list using the load_single_document function, and appends the results to the results list.
5. Returns the list of loaded documents.

loaded_documents = load_documents(git_dir+'source_documents')
print(f"Length of loaded documents: {len(loaded_documents)}")
loaded_documents[0]
Loading new documents: 100%|█████████████████████| 1/1 [00:00<00:00, 246.69it/s]
Length of loaded documents: 1
Document(page_content='Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.  \n\nLast year COVID-19 kept us apart. This year we are finally together again. \n\nTonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \n\nWith a duty to one another to the American people to the Constitution. \n\nAnd with an unwavering resolve that freedom will always triumph over tyranny. \n\nSix days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. \n\nHe thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. \n\nHe met the Ukrainian people. \n\nFrom President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world. \n\nGroups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. \n\nIn this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight. \n\nLet each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world. \n\nPlease rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people. \n\nThroughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos.   \n\nThey keep moving.   \n\nAnd the costs and the threats to America and the world keep rising.   \n\nThat’s why the NATO Alliance was created to secure peace and stability in Europe after World War 2. \n\nThe United States is a member along with 29 other nations. \n\nIt matters. American diplomacy matters. American resolve matters. \n\nPutin’s latest attack on Ukraine was premeditated and unprovoked. \n\nHe rejected repeated efforts at diplomacy. \n\nHe thought the West and NATO wouldn’t respond. And he thought he could divide us at home. Putin was wrong. We were ready.  Here is what we did.   \n\nWe prepared extensively and carefully. \n\nWe spent months building a coalition of other freedom-loving nations from Europe and the Americas to Asia and Africa to confront Putin. \n\nI spent countless hours unifying our European allies. We shared with the world in advance what we knew Putin was planning and precisely how he would try to falsely justify his aggression.  \n\nWe countered Russia’s lies with truth.   \n\nAnd now that he has acted the free world is holding him accountable. \n\nAlong with twenty-seven members of the European Union including France, Germany, Italy, as well as countries like the United Kingdom, Canada, Japan, Korea, Australia, New Zealand, and many others, even Switzerland. \n\nWe are inflicting pain on Russia and supporting the people of Ukraine. Putin is now isolated from the world more than ever. \n\nTogether with our allies –we are right now enforcing powerful economic sanctions. \n\nWe are cutting off Russia’s largest banks from the international financial system.  \n\nPreventing Russia’s central bank from defending the Russian Ruble making Putin’s $630 Billion “war fund” worthless.   \n\nWe are choking off Russia’s access to technology that will sap its economic strength and weaken its military for years to come.  \n\nTonight I say to the Russian oligarchs and corrupt leaders who have bilked billions of dollars off this violent regime no more. \n\nThe U.S. Department of Justice is assembling a dedicated task force to go after the crimes of Russian oligarchs.  \n\nWe are joining with our European allies to find and seize your yachts your luxury apartments your private jets. We are coming for your ill-begotten gains. \n\nAnd tonight I am announcing that we will join our allies in closing off American air space to all Russian flights – further isolating Russia – and adding an additional squeeze –on their economy. The Ruble has lost 30% of its value. \n\nThe Russian stock market has lost 40% of its value and trading remains suspended. Russia’s economy is reeling and Putin alone is to blame. \n\nTogether with our allies we are providing support to the Ukrainians in their fight for freedom. Military assistance. Economic assistance. Humanitarian assistance. \n\nWe are giving more than $1 Billion in direct assistance to Ukraine. \n\nAnd we will continue to aid the Ukrainian people as they defend their country and to help ease their suffering.  \n\nLet me be clear, our forces are not engaged and will not engage in conflict with Russian forces in Ukraine.  \n\nOur forces are not going to Europe to fight in Ukraine, but to defend our NATO Allies – in the event that Putin decides to keep moving west.  \n\nFor that purpose we’ve mobilized American ground forces, air squadrons, and ship deployments to protect NATO countries including Poland, Romania, Latvia, Lithuania, and Estonia. \n\nAs I have made crystal clear the United States and our Allies will defend every inch of territory of NATO countries with the full force of our collective power.  \n\nAnd we remain clear-eyed. The Ukrainians are fighting back with pure courage. But the next few days weeks, months, will be hard on them.  \n\nPutin has unleashed violence and chaos.  But while he may make gains on the battlefield – he will pay a continuing high price over the long run. \n\nAnd a proud Ukrainian people, who have known 30 years  of independence, have repeatedly shown that they will not tolerate anyone who tries to take their country backwards.  \n\nTo all Americans, I will be honest with you, as I’ve always promised. A Russian dictator, invading a foreign country, has costs around the world. \n\nAnd I’m taking robust action to make sure the pain of our sanctions  is targeted at Russia’s economy. And I will use every tool at our disposal to protect American businesses and consumers. \n\nTonight, I can announce that the United States has worked with 30 other countries to release 60 Million barrels of oil from reserves around the world.  \n\nAmerica will lead that effort, releasing 30 Million barrels from our own Strategic Petroleum Reserve. And we stand ready to do more if necessary, unified with our allies.  \n\nThese steps will help blunt gas prices here at home. And I know the news about what’s happening can seem alarming. \n\nBut I want you to know that we are going to be okay. \n\nWhen the history of this era is written Putin’s war on Ukraine will have left Russia weaker and the rest of the world stronger. \n\nWhile it shouldn’t have taken something so terrible for people around the world to see what’s at stake now everyone sees it clearly. \n\nWe see the unity among leaders of nations and a more unified Europe a more unified West. And we see unity among the people who are gathering in cities in large crowds around the world even in Russia to demonstrate their support for Ukraine.  \n\nIn the battle between democracy and autocracy, democracies are rising to the moment, and the world is clearly choosing the side of peace and security. \n\nThis is a real test. It’s going to take time. So let us continue to draw inspiration from the iron will of the Ukrainian people. \n\nTo our fellow Ukrainian Americans who forge a deep bond that connects our two nations we stand with you. \n\nPutin may circle Kyiv with tanks, but he will never gain the hearts and souls of the Ukrainian people. \n\nHe will never extinguish their love of freedom. He will never weaken the resolve of the free world. \n\nWe meet tonight in an America that has lived through two of the hardest years this nation has ever faced. \n\nThe pandemic has been punishing. \n\nAnd so many families are living paycheck to paycheck, struggling to keep up with the rising cost of food, gas, housing, and so much more. \n\nI understand. \n\nI remember when my Dad had to leave our home in Scranton, Pennsylvania to find work. I grew up in a family where if the price of food went up, you felt it. \n\nThat’s why one of the first things I did as President was fight to pass the American Rescue Plan.  \n\nBecause people were hurting. We needed to act, and we did. \n\nFew pieces of legislation have done more in a critical moment in our history to lift us out of crisis. \n\nIt fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans.  \n\nHelped put food on their table, keep a roof over their heads, and cut the cost of health insurance. \n\nAnd as my Dad used to say, it gave people a little breathing room. \n\nAnd unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind. \n\nAnd it worked. It created jobs. Lots of jobs. \n\nIn fact—our economy created over 6.5 Million new jobs just last year, more jobs created in one year  \nthan ever before in the history of America. \n\nOur economy grew at a rate of 5.7% last year, the strongest growth in nearly 40 years, the first step in bringing fundamental change to an economy that hasn’t worked for the working people of this nation for too long.  \n\nFor the past 40 years we were told that if we gave tax breaks to those at the very top, the benefits would trickle down to everyone else. \n\nBut that trickle-down theory led to weaker economic growth, lower wages, bigger deficits, and the widest gap between those at the top and everyone else in nearly a century. \n\nVice President Harris and I ran for office with a new economic vision for America. \n\nInvest in America. Educate Americans. Grow the workforce. Build the economy from the bottom up  \nand the middle out, not from the top down.  \n\nBecause we know that when the middle class grows, the poor have a ladder up and the wealthy do very well. \n\nAmerica used to have the best roads, bridges, and airports on Earth. \n\nNow our infrastructure is ranked 13th in the world. \n\nWe won’t be able to compete for the jobs of the 21st Century if we don’t fix that. \n\nThat’s why it was so important to pass the Bipartisan Infrastructure Law—the most sweeping investment to rebuild America in history. \n\nThis was a bipartisan effort, and I want to thank the members of both parties who worked to make it happen. \n\nWe’re done talking about infrastructure weeks. \n\nWe’re going to have an infrastructure decade. \n\nIt is going to transform America and put us on a path to win the economic competition of the 21st Century that we face with the rest of the world—particularly with China.  \n\nAs I’ve told Xi Jinping, it is never a good bet to bet against the American people. \n\nWe’ll create good jobs for millions of Americans, modernizing roads, airports, ports, and waterways all across America. \n\nAnd we’ll do it all to withstand the devastating effects of the climate crisis and promote environmental justice. \n\nWe’ll build a national network of 500,000 electric vehicle charging stations, begin to replace poisonous lead pipes—so every child—and every American—has clean water to drink at home and at school, provide affordable high-speed internet for every American—urban, suburban, rural, and tribal communities. \n\n4,000 projects have already been announced. \n\nAnd tonight, I’m announcing that this year we will start fixing over 65,000 miles of highway and 1,500 bridges in disrepair. \n\nWhen we use taxpayer dollars to rebuild America – we are going to Buy American: buy American products to support American jobs. \n\nThe federal government spends about $600 Billion a year to keep the country safe and secure. \n\nThere’s been a law on the books for almost a century \nto make sure taxpayers’ dollars support American jobs and businesses. \n\nEvery Administration says they’ll do it, but we are actually doing it. \n\nWe will buy American to make sure everything from the deck of an aircraft carrier to the steel on highway guardrails are made in America. \n\nBut to compete for the best jobs of the future, we also need to level the playing field with China and other competitors. \n\nThat’s why it is so important to pass the Bipartisan Innovation Act sitting in Congress that will make record investments in emerging technologies and American manufacturing. \n\nLet me give you one example of why it’s so important to pass it. \n\nIf you travel 20 miles east of Columbus, Ohio, you’ll find 1,000 empty acres of land. \n\nIt won’t look like much, but if you stop and look closely, you’ll see a “Field of dreams,” the ground on which America’s future will be built. \n\nThis is where Intel, the American company that helped build Silicon Valley, is going to build its $20 billion semiconductor “mega site”. \n\nUp to eight state-of-the-art factories in one place. 10,000 new good-paying jobs. \n\nSome of the most sophisticated manufacturing in the world to make computer chips the size of a fingertip that power the world and our everyday lives. \n\nSmartphones. The Internet. Technology we have yet to invent. \n\nBut that’s just the beginning. \n\nIntel’s CEO, Pat Gelsinger, who is here tonight, told me they are ready to increase their investment from  \n$20 billion to $100 billion. \n\nThat would be one of the biggest investments in manufacturing in American history. \n\nAnd all they’re waiting for is for you to pass this bill. \n\nSo let’s not wait any longer. Send it to my desk. I’ll sign it.  \n\nAnd we will really take off. \n\nAnd Intel is not alone. \n\nThere’s something happening in America. \n\nJust look around and you’ll see an amazing story. \n\nThe rebirth of the pride that comes from stamping products “Made In America.” The revitalization of American manufacturing.   \n\nCompanies are choosing to build new factories here, when just a few years ago, they would have built them overseas. \n\nThat’s what is happening. Ford is investing $11 billion to build electric vehicles, creating 11,000 jobs across the country. \n\nGM is making the largest investment in its history—$7 billion to build electric vehicles, creating 4,000 jobs in Michigan. \n\nAll told, we created 369,000 new manufacturing jobs in America just last year. \n\nPowered by people I’ve met like JoJo Burgess, from generations of union steelworkers from Pittsburgh, who’s here with us tonight. \n\nAs Ohio Senator Sherrod Brown says, “It’s time to bury the label “Rust Belt.” \n\nIt’s time. \n\nBut with all the bright spots in our economy, record job growth and higher wages, too many families are struggling to keep up with the bills.  \n\nInflation is robbing them of the gains they might otherwise feel. \n\nI get it. That’s why my top priority is getting prices under control. \n\nLook, our economy roared back faster than most predicted, but the pandemic meant that businesses had a hard time hiring enough workers to keep up production in their factories. \n\nThe pandemic also disrupted global supply chains. \n\nWhen factories close, it takes longer to make goods and get them from the warehouse to the store, and prices go up. \n\nLook at cars. \n\nLast year, there weren’t enough semiconductors to make all the cars that people wanted to buy. \n\nAnd guess what, prices of automobiles went up. \n\nSo—we have a choice. \n\nOne way to fight inflation is to drive down wages and make Americans poorer.  \n\nI have a better plan to fight inflation. \n\nLower your costs, not your wages. \n\nMake more cars and semiconductors in America. \n\nMore infrastructure and innovation in America. \n\nMore goods moving faster and cheaper in America. \n\nMore jobs where you can earn a good living in America. \n\nAnd instead of relying on foreign supply chains, let’s make it in America. \n\nEconomists call it “increasing the productive capacity of our economy.” \n\nI call it building a better America. \n\nMy plan to fight inflation will lower your costs and lower the deficit. \n\n17 Nobel laureates in economics say my plan will ease long-term inflationary pressures. Top business leaders and most Americans support my plan. And here’s the plan: \n\nFirst – cut the cost of prescription drugs. Just look at insulin. One in ten Americans has diabetes. In Virginia, I met a 13-year-old boy named Joshua Davis.  \n\nHe and his Dad both have Type 1 diabetes, which means they need insulin every day. Insulin costs about $10 a vial to make.  \n\nBut drug companies charge families like Joshua and his Dad up to 30 times more. I spoke with Joshua’s mom. \n\nImagine what it’s like to look at your child who needs insulin and have no idea how you’re going to pay for it.  \n\nWhat it does to your dignity, your ability to look your child in the eye, to be the parent you expect to be. \n\nJoshua is here with us tonight. Yesterday was his birthday. Happy birthday, buddy.  \n\nFor Joshua, and for the 200,000 other young people with Type 1 diabetes, let’s cap the cost of insulin at $35 a month so everyone can afford it.  \n\nDrug companies will still do very well. And while we’re at it let Medicare negotiate lower prices for prescription drugs, like the VA already does. \n\nLook, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. \n\nSecond – cut energy costs for families an average of $500 a year by combatting climate change.  \n\nLet’s provide investments and tax credits to weatherize your homes and businesses to be energy efficient and you get a tax credit; double America’s clean energy production in solar, wind, and so much more;  lower the price of electric vehicles, saving you another $80 a month because you’ll never have to pay at the gas pump again. \n\nThird – cut the cost of child care. Many families pay up to $14,000 a year for child care per child.  \n\nMiddle-class and working families shouldn’t have to pay more than 7% of their income for care of young children.  \n\nMy plan will cut the cost in half for most families and help parents, including millions of women, who left the workforce during the pandemic because they couldn’t afford child care, to be able to get back to work. \n\nMy plan doesn’t stop there. It also includes home and long-term care. More affordable housing. And Pre-K for every 3- and 4-year-old.  \n\nAll of these will lower costs. \n\nAnd under my plan, nobody earning less than $400,000 a year will pay an additional penny in new taxes. Nobody.  \n\nThe one thing all Americans agree on is that the tax system is not fair. We have to fix it.  \n\nI’m not looking to punish anyone. But let’s make sure corporations and the wealthiest Americans start paying their fair share. \n\nJust last year, 55 Fortune 500 corporations earned $40 billion in profits and paid zero dollars in federal income tax.  \n\nThat’s simply not fair. That’s why I’ve proposed a 15% minimum tax rate for corporations. \n\nWe got more than 130 countries to agree on a global minimum tax rate so companies can’t get out of paying their taxes at home by shipping jobs and factories overseas. \n\nThat’s why I’ve proposed closing loopholes so the very wealthy don’t pay a lower tax rate than a teacher or a firefighter.  \n\nSo that’s my plan. It will grow the economy and lower costs for families. \n\nSo what are we waiting for? Let’s get this done. And while you’re at it, confirm my nominees to the Federal Reserve, which plays a critical role in fighting inflation.  \n\nMy plan will not only lower costs to give families a fair shot, it will lower the deficit. \n\nThe previous Administration not only ballooned the deficit with tax cuts for the very wealthy and corporations, it undermined the watchdogs whose job was to keep pandemic relief funds from being wasted. \n\nBut in my administration, the watchdogs have been welcomed back. \n\nWe’re going after the criminals who stole billions in relief money meant for small businesses and millions of Americans.  \n\nAnd tonight, I’m announcing that the Justice Department will name a chief prosecutor for pandemic fraud. \n\nBy the end of this year, the deficit will be down to less than half what it was before I took office.  \n\nThe only president ever to cut the deficit by more than one trillion dollars in a single year. \n\nLowering your costs also means demanding more competition. \n\nI’m a capitalist, but capitalism without competition isn’t capitalism. \n\nIt’s exploitation—and it drives up prices. \n\nWhen corporations don’t have to compete, their profits go up, your prices go up, and small businesses and family farmers and ranchers go under. \n\nWe see it happening with ocean carriers moving goods in and out of America. \n\nDuring the pandemic, these foreign-owned companies raised prices by as much as 1,000% and made record profits. \n\nTonight, I’m announcing a crackdown on these companies overcharging American businesses and consumers. \n\nAnd as Wall Street firms take over more nursing homes, quality in those homes has gone down and costs have gone up.  \n\nThat ends on my watch. \n\nMedicare is going to set higher standards for nursing homes and make sure your loved ones get the care they deserve and expect. \n\nWe’ll also cut costs and keep the economy going strong by giving workers a fair shot, provide more training and apprenticeships, hire them based on their skills not degrees. \n\nLet’s pass the Paycheck Fairness Act and paid leave.  \n\nRaise the minimum wage to $15 an hour and extend the Child Tax Credit, so no one has to raise a family in poverty. \n\nLet’s increase Pell Grants and increase our historic support of HBCUs, and invest in what Jill—our First Lady who teaches full-time—calls America’s best-kept secret: community colleges. \n\nAnd let’s pass the PRO Act when a majority of workers want to form a union—they shouldn’t be stopped.  \n\nWhen we invest in our workers, when we build the economy from the bottom up and the middle out together, we can do something we haven’t done in a long time: build a better America. \n\nFor more than two years, COVID-19 has impacted every decision in our lives and the life of the nation. \n\nAnd I know you’re tired, frustrated, and exhausted. \n\nBut I also know this. \n\nBecause of the progress we’ve made, because of your resilience and the tools we have, tonight I can say  \nwe are moving forward safely, back to more normal routines.  \n\nWe’ve reached a new moment in the fight against COVID-19, with severe cases down to a level not seen since last July.  \n\nJust a few days ago, the Centers for Disease Control and Prevention—the CDC—issued new mask guidelines. \n\nUnder these new guidelines, most Americans in most of the country can now be mask free.   \n\nAnd based on the projections, more of the country will reach that point across the next couple of weeks. \n\nThanks to the progress we have made this past year, COVID-19 need no longer control our lives.  \n\nI know some are talking about “living with COVID-19”. Tonight – I say that we will never just accept living with COVID-19. \n\nWe will continue to combat the virus as we do other diseases. And because this is a virus that mutates and spreads, we will stay on guard. \n\nHere are four common sense steps as we move forward safely.  \n\nFirst, stay protected with vaccines and treatments. We know how incredibly effective vaccines are. If you’re vaccinated and boosted you have the highest degree of protection. \n\nWe will never give up on vaccinating more Americans. Now, I know parents with kids under 5 are eager to see a vaccine authorized for their children. \n\nThe scientists are working hard to get that done and we’ll be ready with plenty of vaccines when they do. \n\nWe’re also ready with anti-viral treatments. If you get COVID-19, the Pfizer pill reduces your chances of ending up in the hospital by 90%.  \n\nWe’ve ordered more of these pills than anyone in the world. And Pfizer is working overtime to get us 1 Million pills this month and more than double that next month.  \n\nAnd we’re launching the “Test to Treat” initiative so people can get tested at a pharmacy, and if they’re positive, receive antiviral pills on the spot at no cost.  \n\nIf you’re immunocompromised or have some other vulnerability, we have treatments and free high-quality masks. \n\nWe’re leaving no one behind or ignoring anyone’s needs as we move forward. \n\nAnd on testing, we have made hundreds of millions of tests available for you to order for free.   \n\nEven if you already ordered free tests tonight, I am announcing that you can order more from covidtests.gov starting next week. \n\nSecond – we must prepare for new variants. Over the past year, we’ve gotten much better at detecting new variants. \n\nIf necessary, we’ll be able to deploy new vaccines within 100 days instead of many more months or years.  \n\nAnd, if Congress provides the funds we need, we’ll have new stockpiles of tests, masks, and pills ready if needed. \n\nI cannot promise a new variant won’t come. But I can promise you we’ll do everything within our power to be ready if it does.  \n\nThird – we can end the shutdown of schools and businesses. We have the tools we need. \n\nIt’s time for Americans to get back to work and fill our great downtowns again.  People working from home can feel safe to begin to return to the office.   \n\nWe’re doing that here in the federal government. The vast majority of federal workers will once again work in person. \n\nOur schools are open. Let’s keep it that way. Our kids need to be in school. \n\nAnd with 75% of adult Americans fully vaccinated and hospitalizations down by 77%, most Americans can remove their masks, return to work, stay in the classroom, and move forward safely. \n\nWe achieved this because we provided free vaccines, treatments, tests, and masks. \n\nOf course, continuing this costs money. \n\nI will soon send Congress a request. \n\nThe vast majority of Americans have used these tools and may want to again, so I expect Congress to pass it quickly.   \n\nFourth, we will continue vaccinating the world.     \n\nWe’ve sent 475 Million vaccine doses to 112 countries, more than any other nation. \n\nAnd we won’t stop. \n\nWe have lost so much to COVID-19. Time with one another. And worst of all, so much loss of life. \n\nLet’s use this moment to reset. Let’s stop looking at COVID-19 as a partisan dividing line and see it for what it is: A God-awful disease.  \n\nLet’s stop seeing each other as enemies, and start seeing each other for who we really are: Fellow Americans.  \n\nWe can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. \n\nI recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. \n\nThey were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. \n\nOfficer Mora was 27 years old. \n\nOfficer Rivera was 22. \n\nBoth Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. \n\nI spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \n\nI’ve worked on these issues a long time. \n\nI know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. \n\nSo let’s not abandon our streets. Or choose between safety and equal justice. \n\nLet’s come together to protect our communities, restore trust, and hold law enforcement accountable. \n\nThat’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. \n\nThat’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.  \n\nWe should all agree: The answer is not to Defund the police. The answer is to FUND the police with the resources and training they need to protect our communities. \n\nI ask Democrats and Republicans alike: Pass my budget and keep our neighborhoods safe.  \n\nAnd I will keep doing everything in my power to crack down on gun trafficking and ghost guns you can buy online and make at home—they have no serial numbers and can’t be traced. \n\nAnd I ask Congress to pass proven measures to reduce gun violence. Pass universal background checks. Why should anyone on a terrorist list be able to purchase a weapon? \n\nBan assault weapons and high-capacity magazines. \n\nRepeal the liability shield that makes gun manufacturers the only industry in America that can’t be sued. \n\nThese laws don’t infringe on the Second Amendment. They save lives. \n\nThe most fundamental right in America is the right to vote – and to have it counted. And it’s under assault. \n\nIn state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. \n\nWe cannot let this happen. \n\nTonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n\nTonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n\nOne of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n\nAnd I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. \n\nA former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling.  \n\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers.  \n\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders. \n\nWe can do all this while keeping lit the torch of liberty that has led generations of immigrants to this land—my forefathers and so many of yours. \n\nProvide a pathway to citizenship for Dreamers, those on temporary status, farm workers, and essential workers. \n\nRevise our laws so businesses have the workers they need and families don’t wait decades to reunite. \n\nIt’s not only the right thing to do—it’s the economically smart thing to do. \n\nThat’s why immigration reform is supported by everyone from labor unions to religious leaders to the U.S. Chamber of Commerce. \n\nLet’s get it done once and for all. \n\nAdvancing liberty and justice also requires protecting the rights of women. \n\nThe constitutional right affirmed in Roe v. Wade—standing precedent for half a century—is under attack as never before. \n\nIf we want to go forward—not backward—we must protect access to health care. Preserve a woman’s right to choose. And let’s continue to advance maternal health care in America. \n\nAnd for our LGBTQ+ Americans, let’s finally get the bipartisan Equality Act to my desk. The onslaught of state laws targeting transgender Americans and their families is wrong. \n\nAs I said last year, especially to our younger transgender Americans, I will always have your back as your President, so you can be yourself and reach your God-given potential. \n\nWhile it often appears that we never agree, that isn’t true. I signed 80 bipartisan bills into law last year. From preventing government shutdowns to protecting Asian-Americans from still-too-common hate crimes to reforming military justice. \n\nAnd soon, we’ll strengthen the Violence Against Women Act that I first wrote three decades ago. It is important for us to show the nation that we can come together and do big things. \n\nSo tonight I’m offering a Unity Agenda for the Nation. Four big things we can do together.  \n\nFirst, beat the opioid epidemic. \n\nThere is so much we can do. Increase funding for prevention, treatment, harm reduction, and recovery.  \n\nGet rid of outdated rules that stop doctors from prescribing treatments. And stop the flow of illicit drugs by working with state and local law enforcement to go after traffickers. \n\nIf you’re suffering from addiction, know you are not alone. I believe in recovery, and I celebrate the 23 million Americans in recovery. \n\nSecond, let’s take on mental health. Especially among our children, whose lives and education have been turned upside down.  \n\nThe American Rescue Plan gave schools money to hire teachers and help students make up for lost learning.  \n\nI urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. \n\nChildren were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media. \n\nAs Frances Haugen, who is here with us tonight, has shown, we must hold social media platforms accountable for the national experiment they’re conducting on our children for profit. \n\nIt’s time to strengthen privacy protections, ban targeted advertising to children, demand tech companies stop collecting personal data on our children. \n\nAnd let’s get all Americans the mental health services they need. More people they can turn to for help, and full parity between physical and mental health care. \n\nThird, support our veterans. \n\nVeterans are the best of us. \n\nI’ve always believed that we have a sacred obligation to equip all those we send to war and care for them and their families when they come home. \n\nMy administration is providing assistance with job training and housing, and now helping lower-income veterans get VA care debt-free.  \n\nOur troops in Iraq and Afghanistan faced many dangers. \n\nOne was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. \n\nWhen they came home, many of the world’s fittest and best trained warriors were never the same. \n\nHeadaches. Numbness. Dizziness. \n\nA cancer that would put them in a flag-draped coffin. \n\nI know. \n\nOne of those soldiers was my son Major Beau Biden. \n\nWe don’t know for sure if a burn pit was the cause of his brain cancer, or the diseases of so many of our troops. \n\nBut I’m committed to finding out everything we can. \n\nCommitted to military families like Danielle Robinson from Ohio. \n\nThe widow of Sergeant First Class Heath Robinson.  \n\nHe was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. \n\nStationed near Baghdad, just yards from burn pits the size of football fields. \n\nHeath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter. \n\nBut cancer from prolonged exposure to burn pits ravaged Heath’s lungs and body. \n\nDanielle says Heath was a fighter to the very end. \n\nHe didn’t know how to stop fighting, and neither did she. \n\nThrough her pain she found purpose to demand we do better. \n\nTonight, Danielle—we are. \n\nThe VA is pioneering new ways of linking toxic exposures to diseases, already helping more veterans get benefits. \n\nAnd tonight, I’m announcing we’re expanding eligibility to veterans suffering from nine respiratory cancers. \n\nI’m also calling on Congress: pass a law to make sure veterans devastated by toxic exposures in Iraq and Afghanistan finally get the benefits and comprehensive health care they deserve. \n\nAnd fourth, let’s end cancer as we know it. \n\nThis is personal to me and Jill, to Kamala, and to so many of you. \n\nCancer is the #2 cause of death in America–second only to heart disease. \n\nLast month, I announced our plan to supercharge  \nthe Cancer Moonshot that President Obama asked me to lead six years ago. \n\nOur goal is to cut the cancer death rate by at least 50% over the next 25 years, turn more cancers from death sentences into treatable diseases.  \n\nMore support for patients and families. \n\nTo get there, I call on Congress to fund ARPA-H, the Advanced Research Projects Agency for Health. \n\nIt’s based on DARPA—the Defense Department project that led to the Internet, GPS, and so much more.  \n\nARPA-H will have a singular purpose—to drive breakthroughs in cancer, Alzheimer’s, diabetes, and more. \n\nA unity agenda for the nation. \n\nWe can do this. \n\nMy fellow Americans—tonight , we have gathered in a sacred space—the citadel of our democracy. \n\nIn this Capitol, generation after generation, Americans have debated great questions amid great strife, and have done great things. \n\nWe have fought for freedom, expanded liberty, defeated totalitarianism and terror. \n\nAnd built the strongest, freest, and most prosperous nation the world has ever known. \n\nNow is the hour. \n\nOur moment of responsibility. \n\nOur test of resolve and conscience, of history itself. \n\nIt is in this moment that our character is formed. Our purpose is found. Our future is forged. \n\nWell I know this nation.  \n\nWe will meet the test. \n\nTo protect freedom and liberty, to expand fairness and opportunity. \n\nWe will save democracy. \n\nAs hard as these times have been, I am more optimistic about America today than I have been my whole life. \n\nBecause I see the future that is within our grasp. \n\nBecause I know there is simply nothing beyond our capacity. \n\nWe are the only nation on Earth that has always turned every crisis we have faced into an opportunity. \n\nThe only nation that can be defined by a single word: possibilities. \n\nSo on this night, in our 245th year as a nation, I have come to report on the State of the Union. \n\nAnd my report is this: the State of the Union is strong—because you, the American people, are strong. \n\nWe are stronger today than we were a year ago. \n\nAnd we will be stronger a year from now than we are today. \n\nNow is our moment to meet and overcome the challenges of our time. \n\nAnd we will, as one people. \n\nOne America. \n\nThe United States of America. \n\nMay God bless you all. May God protect our troops.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'})

You can see we have loaded the state_of_the_union.txt file from the privateGPT repo. As this is the only file in that directory the length of loaded documents is one.

3.1.2 Splitting the documents into smaller chunks

Now we have seen how we can load multiple documents of different extensions using the load_documents function. The next step is to look at process_document function which loads and splits large documents into smaller chunks.

from langchain.text_splitter import RecursiveCharacterTextSplitter

chunk_size = 500
chunk_overlap = 50

def process_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
    """
    Load documents and split in chunks
    """
    print(f"Loading documents from {source_dir}")
    documents = load_documents(source_dir, ignored_files)
    if not documents:
        print("No new documents to load")
        exit(0)
    print(f"Loaded {len(documents)} new documents from {source_dir}")
    ## Load text splitter
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    ## Split text
    texts = text_splitter.split_documents(documents)
    print(f"Split into {len(texts)} chunks of text (max. {chunk_size} tokens each)")
    return texts

processed_documents = process_documents(git_dir+'source_documents')
Loading documents from ../../../../privateGPT/source_documents
Loading new documents: 100%|█████████████████████| 1/1 [00:00<00:00, 315.74it/s]
Loaded 1 new documents from ../../../../privateGPT/source_documents
Split into 90 chunks of text (max. 500 tokens each)

The process_documents function performs the following steps:

  1. Loads all the documents from the source_dir directory using the load_documents function.
  2. Initializes an instance of RecursiveCharacterTextSplitter from the langchain.text_splitter module, providing the chunk_size and chunk_overlap parameters. This class is responsible for splitting a list of documents into smaller overlapping chunks. [RecursiveCharacterTextSplitter documentation].
  3. Uses the split_documents method of the RecursiveCharacterTextSplitter instance to split the loaded documents into smaller chunks.
  4. Returns the resulting list of the smaller document chunks.

3.1.3 Initializing the embedding model

Next, we load our embedding module which converts the smaller document chunks from previous steps to embeddings.

from langchain.embeddings import HuggingFaceEmbeddings
EMBEDDINGS_MODEL_NAME = "all-MiniLM-L6-v2"
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)

print("Testing on a single query.")
embedded_vector = embeddings.embed_query("What is your name?")
print(f"Size of embedded vector: {len(embedded_vector)}")
Testing on a single query.
Size of embedded vector: 384

The given code snippet carries out the following steps:

  1. Imports the HuggingFaceEmbeddings function from the langchain.embeddings module. This function is responsible for loading and encapsulating the SentenceTransformers embeddings, which are used for generating dense vector representations of sentences. You can refer to the HuggingFaceEmbeddings documentation for more details.
  2. Loads the all-MiniLM-L6-v2 model from the sentence_transformers library. This model is specifically designed to map sentences and paragraphs into a 384-dimensional dense vector space. It is commonly utilized for tasks such as semantic search and similarity analysis.

We can see that our embedded vector on a sample query returns a 384 dimension vector.

3.1.4 Embed smaller text and save it in the vector database

The next step involves utilizing the document chunks and the embedding model to store the documents and their corresponding embeddings in a vector database.

from chromadb.config import Settings
from langchain.vectorstores import Chroma

PERSIST_DIRECTORY= git_dir+"db"
# Define the Chroma settings
CHROMA_SETTINGS = Settings(
        chroma_db_impl='duckdb+parquet',
        persist_directory=PERSIST_DIRECTORY,
        anonymized_telemetry=False
)
## Create the embedding database
db = Chroma.from_documents(processed_documents, embeddings, persist_directory=PERSIST_DIRECTORY, client_settings=CHROMA_SETTINGS)
db.persist()
Using embedded DuckDB with persistence: data will be stored in: ../../../../privateGPT/db

The given code snippet performs the following operations:

  1. It imports the Settings class from the chromadb.config module and the Chroma class from the langchain.vectorstores module.
  2. It creates an instance of the Settings class named CHROMA_SETTINGS, providing several configuration parameters:
    • chroma_db_impl is set to 'duckdb+parquet', specifying the implementation to be used for the Chroma vector database.
    • persist_directory is set to the PERSIST_DIRECTORY variable defined earlier, indicating the directory where the vector database will be saved.
    • anonymized_telemetry is set to False, indicating whether anonymized telemetry data should be collected.
  3. It creates a vector database by calling the Chroma.from_documents() method. This method takes the following arguments:
    • processed_documents: The list of processed documents obtained from the previous step.
    • embeddings: The embeddings object/model used to generate the document embeddings.
    • persist_directory: The directory where the vector database will be persisted, specified by the PERSIST_DIRECTORY variable.
    • client_settings: The settings object (CHROMA_SETTINGS) containing configuration parameters for the vector database.
  4. We use db.persist() to store the index for future retrieval task
## Test the semantic retrieval 
db.similarity_search(query="What is the American Rescue Plan?", k= 4)
[Document(page_content='The American Rescue Plan gave schools money to hire teachers and help students make up for lost learning.  \n\nI urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. \n\nChildren were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'}),
 Document(page_content='It fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans.  \n\nHelped put food on their table, keep a roof over their heads, and cut the cost of health insurance. \n\nAnd as my Dad used to say, it gave people a little breathing room. \n\nAnd unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'}),
 Document(page_content='Look, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. \n\nSecond – cut energy costs for families an average of $500 a year by combatting climate change.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'}),
 Document(page_content='That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. \n\nThat’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'})]
db = None

To test the retrieval of semantic similarity, we can use the similarity_search function. similarity_search function takes a text query as input and returns the top k=4 document chunks from the vector database.

3.2 Question & Answer Interface

Let’s explore the Q&A interface in more detail. The Q&A interface consists of the following steps:

  1. Load the vector database and prepare it for the retrieval task.
  2. Load a pre-trained Large language model from LlamaCpp or GPT4ALL.
  3. Prompt the user with a query and generate a response using the RetrievalQA pipeline from langchain.chains.
Fig.6: Question Answering Pipeline

Let’s look at these steps one by one.

3.2.1 Load the vector database

First, we import the required libraries.

from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from chromadb.config import Settings
git_dir = "../../../../privateGPT/"
PERSIST_DIRECTORY= git_dir+"db"
EMBEDDINGS_MODEL_NAME = "all-MiniLM-L6-v2"

# Define the Chroma settings
CHROMA_SETTINGS = Settings(
        chroma_db_impl='duckdb+parquet',
        persist_directory=PERSIST_DIRECTORY,
        anonymized_telemetry=False
)

embeddings = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_NAME)
db = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
retriever = db.as_retriever()
Using embedded DuckDB with persistence: data will be stored in: ../../../../privateGPT/db

The given code snippet carries out the following steps:

  1. Loads the embeddings using the HuggingFaceEmbeddings function, which was previously used to create the embedding store.
  2. Instantiates a Chroma vector database that was created earlier.
  3. Sets the vector database in retrieval mode.
## Testing retriever
retriever.vectorstore.similarity_search(query = "What is Amercian rescue plan?")
[Document(page_content='The American Rescue Plan gave schools money to hire teachers and help students make up for lost learning.  \n\nI urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. \n\nChildren were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'}),
 Document(page_content='It fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans.  \n\nHelped put food on their table, keep a roof over their heads, and cut the cost of health insurance. \n\nAnd as my Dad used to say, it gave people a little breathing room. \n\nAnd unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'}),
 Document(page_content='That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. \n\nThat’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'}),
 Document(page_content='Look, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. \n\nSecond – cut energy costs for families an average of $500 a year by combatting climate change.', metadata={'source': '../../../../privateGPT/source_documents/state_of_the_union.txt'})]

3.2.2 Load a pre-trained Large language model.

from langchain.llms import GPT4All

MODEL_PATH = git_dir+"models/ggml-gpt4all-j-v1.3-groovy.bin" 
MODEL_N_CTX=1000

# Prepare the LLM
llm = GPT4All(model=MODEL_PATH, n_ctx=MODEL_N_CTX, backend='gptj', callbacks=None, verbose=False)
gptj_model_load: loading model from '../../../../privateGPT/models/ggml-gpt4all-j-v1.3-groovy.bin' - please wait ...
gptj_model_load: n_vocab = 50400
gptj_model_load: n_ctx   = 2048
gptj_model_load: n_embd  = 4096
gptj_model_load: n_head  = 16
gptj_model_load: n_layer = 28
gptj_model_load: n_rot   = 64
gptj_model_load: f16     = 2
gptj_model_load: ggml ctx size = 4505.45 MB
gptj_model_load: memory_size =   896.00 MB, n_mem = 57344
gptj_model_load: ................................... done
gptj_model_load: model size =  3609.38 MB / num tensors = 285

The code snippet above create an instance of the GPT4All class named llm, which represents the Language Model (LLM) using the GPT-4All model. The constructor of GPT4All takes the following arguments:
- model: The path to the GPT-4All model file specified by the MODEL_PATH variable.
- n_ctx: The context size or maximum length of input sequences specified by the MODEL_N_CTX variable.
- backend: The backend to use for the LLM. In this case, it is set to ‘gptj’.
- callbacks: The callbacks to be used during the LLM execution. In this case, it is set to None.
- verbose: A boolean flag indicating whether to print verbose output during LLM execution. In this case, it is set to False.

3.2.3 Prompt the user with a query and generate a response

from langchain.chains import RetrievalQA
qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
query = "What is American rescue plan?"
res = qa(query)
answer, docs = res['result'], res['source_documents']

# Get the answer from the chain
# Print the result
print("\n\n> Question:")
print(query)
print("\n> Answer:")
print(answer)

# Print the relevant sources used for the answer
for document in docs:
    print("\n> " + document.metadata["source"] + ":")
    print(document.page_content)


> Question:
What is American rescue plan?

> Answer:
 The American Rescue Plan is a program that provides funding to schools to hire teachers and help students make up for lost learning due to the COVID-19 pandemic. It also provides economic relief for tens of millions of Americans by helping them put food on their table, keep a roof over their heads, and cut the cost of health insurance. The plan also helps working people by providing breathing room and giving them a little breathing room. It is a program that helps millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums and combat climate change by cutting energy costs for families an average of $500 a year.

> ../../../../privateGPT/source_documents/state_of_the_union.txt:
The American Rescue Plan gave schools money to hire teachers and help students make up for lost learning.  

I urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. 

Children were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media.

> ../../../../privateGPT/source_documents/state_of_the_union.txt:
It fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans.  

Helped put food on their table, keep a roof over their heads, and cut the cost of health insurance. 

And as my Dad used to say, it gave people a little breathing room. 

And unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind.

> ../../../../privateGPT/source_documents/state_of_the_union.txt:
Look, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. 

Second – cut energy costs for families an average of $500 a year by combatting climate change.

> ../../../../privateGPT/source_documents/state_of_the_union.txt:
That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. 

That’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.

Firstly, an instance of the RetrievalQA class named qa is created using the from_chain_type method. The RetrievalQA class is a chain specifically designed for question-answering tasks over an index. Please refer to the documentation for further details. The from_chain_type method takes the following arguments:

  • llm: The Language Model instance (llm) that was created previously.
  • chain_type: A string representing the type of chain to be used. In this case, it is set to "stuff". There may be other available chain types specific to the question-answering scenario. Please consult the documentation for more information.
  • retriever: An instance of a Chroma database used to retrieve relevant documents for the given query.
  • return_source_documents: A boolean flag indicating whether to return the source documents along with the answer. In this case, it is set to True.

Next, the qa instance is used to process a query. The Language Model (LLM) within the qa instance generates a response that includes the query, the answer, and the source documents used as context for generating the answer.

Finally, the answer and source documents are printed out for display.

3.3 Conclusion

In this blog post, we explored privateGPT, its implementation, and the code walkthrough for its ingestion pipeline and q&A interface. I hope this blog post has been valuable in understanding privateGPT and its implementation. I recommend my readers to try privateGPT on your own knowledge base.

I hope you enjoyed reading it. If there is any feedback on the code or just the blog post, feel free to comment below or reach out on LinkedIn.

]]> NLP Deep Learning LLMs https://aayushmnit.com/posts/2023-05-22-PrivateGPTWalkthrough/privateGPTWalkthrough.html Mon, 22 May 2023 07:00:00 GMT Python OOPs fundamentals Aayush Agrawal https://aayushmnit.com/posts/2022-12-20-PythonFundamentals/Python OOPs Fundamentals.html

An introduction to Object Oriented programming using Python.

Increasingly it’s becoming important for Data professionals to become better at programming and modern programming is centered around Object Oriented programming paradigm. This article helps in explaining some important programming concepts which are mostly language agnostic but we will be using Python in this article.

Object-oriented programming (OOPs) is a programming paradigm that relies on the concept of classes and objects. The basic idea of OOP is to divide a sophisticated program into a number of objects that interact with each other to achieve the desired functionality. There are several advantages of using OOP for data science:

  • Encapsulation: OOPs allow you to wrap data and the methods that operate on that data within a single unit (i.e., an object). This makes it easier to organize and manage your code, as well as protect the data from accidental modification.
  • Code reusability: OOPs allow you to create reusable code by defining classes that can be used to create multiple objects with the same behavior. This makes it easier to develop and maintain your code, as you can reuse existing code instead of writing new code from scratch.
  • Modularity: OOPs allow you to divide your code into smaller, modular units (i.e., objects). This makes it easier to understand and maintain your code, as you can focus on one piece of functionality at a time.
  • Inheritance: OOPs allow you to create a new class that is a modified version of an existing class. This allows you to reuse code from the existing class and override or extend it as needed. This can save time and reduce the amount of code you need to write.

Overall, OOP can help data professionals organize and manage their code more effectively, making it easier to develop and maintain data science projects. Let’s dive into the OOPs concept.

1 What are Objects and Classes?

Classes are the blueprint for defining an Object. While an Object is a collection of data/properties and their behaviors/methods.

For example- Think of a class Bulb that will have a state (On/Off) and methods to turnOn and turnoff the bulb.

class Bulb():
    def __init__(self, onOff=False): self.onOff = onOff    
    def turnOn(self): self.onOff = True
    def turnOff(self): self.onOff = False

Now we can create multiple bulb objects from this Bulb class.

b1 = Bulb(onOff=True)
b2 = Bulb()
print(f"Bulb 1 state is :{b1.onOff}, Bulb 2 state is :{b2.onOff}")
Bulb 1 state is :True, Bulb 2 state is :False

b1 and b2 are objects of the class Bulb. Let’s use the turnOn and turnOff methods to update the bulb properties.

b1.turnOff(); b2.turnOn()
print(f"Bulb 1 state is :{b1.onOff}, Bulb 2 state is :{b2.onOff}")
Bulb 1 state is :False, Bulb 2 state is :True

We can see from the example above, a Bulb object contains the onOff property. Properties are variables that contain information regarding the object of a class and Methods like turnOn and turnOff in our Bulb class are functions that have access to the properties of a class. Methods can accept additional parameters, modify properties and return values.

2 Class and Instance variables

In Python, properties can be defined in two ways -

  • Class Variables - Class variables are shared by all objects of the class. A change in the class variable will change the value of that property in all the objects of the class.
  • Instance Variables - Instance variables are unique to each instance or object of the class. A change in instance variable will change the value of the property in that specific object only.
class Employee:
    # Creating a class variable
    companyName = "Microsoft"
    
    def __init__(self, name):
        # creating an instance variable
        self.name = name
    
e1 = Employee('Aayush')
e2 = Employee('John')

print(f'Name :{e1.name}')
print(f'Company Name: {e1.companyName}')
print(f'Name :{e2.name}')
print(f'Company Name: {e2.companyName}')
Name :Aayush
Company Name: Microsoft
Name :John
Company Name: Microsoft

We can see above, the class variable is defined outside of the initializer and the instance variable is defined inside the initializer.

Employee.companyName = "Amazon"
print(e1.companyName, e2.companyName)
Amazon Amazon

We can see above changing a class variable in the Employee class changes the class variable in all objects of the class. Most of the time we will be using instance variables but knowledge about class variables can come in handy. Let’s look at an interesting use of class variable -

class Employee:
    # Creating a class variable
    companyName = "Microsoft"
    companyEmployees = []
    
    def __init__(self, name):
        # creating an instance variable
        self.name = name
        self.companyEmployees.append(self.name)
    
e1 = Employee('Aayush')
e2 = Employee('John')

print(f'Name :{e1.name}')
print(f'Team Members: {e1.companyEmployees}')
print(f'Name :{e2.name}')
print(f'Company Name: {e2.companyEmployees}')
Name :Aayush
Team Members: ['Aayush', 'John']
Name :John
Company Name: ['Aayush', 'John']

We can see above, we are saving all objects of the Employee class in companyEmployees which is a list shared by all objects of the class Employee.

3 Class, Static and Instance methods

In Python classes, we have three types of methods -

  • Class Methods - Class methods work with class variables and are accessible using the class name rather than its object.
  • Static Methods - Static methods are methods that are usually limited to class only and not their objects. They don’t typically modify or access class and instance variables. They are used as utility functions inside the class and we don’t want the inherited class to modify them.
  • Instance Methods - Instance methods are the most used methods and have access to instance variables within the class. They can also take new parameters to perform desired operations.
class Employee:
    # Creating a class variable
    companyName = "Microsoft"
    companyEmployees = []
    
    def __init__(self, name):
        # creating an instance variable
        self.name = name
        self.companyEmployees.append(self.name)
    
    @classmethod
    def getCompanyName(cls): # This is a class method
         return cls.companyName
    
    @staticmethod
    def plusTwo(x): # This is a static method
        return x+2
    
    def getName(self): # This is an instance method
        return self.name
    
e1 = Employee('Aayush')
print(f"Calling class method. Company name is {e1.getCompanyName()}")
print(f"Calling Static method. {e1.plusTwo(2)}")
print(f"Calling instance method. Employee name is {e1.getName()}")
Calling class method. Company name is Microsoft
Calling Static method. 4
Calling instance method. Employee name is Aayush

We can see above we use the @classmethod decorator to define the class method. cls is used to refer to the class just as self is used to refer to the object of the class. The class method at least takes one argument cls.

Note

We can use any other name instead of cls but cls is used as a convention.

We use @staticmethod decorator to define static class plusTwo. We can see that static methods don’t take any argument like self and cls.

The most commonly used methods are instance methods and they can be defined without a decorator within the class. Just like the class method they take at least one argument which is self by convention.

Note

We can use any other name instead of self but self is used as a convention.

4 Access Modifiers

Access modifiers limit access to the variables and functions of a class. There are three types of access modifiers - public, protected, and private.

4.1 Public Attributes

Public attributes are those methods and properties which can be accessed anywhere inside and outside of the class. By default, all the member variables and functions are public.

class Employee:
    def __init__(self, name):
        self.name = name ## Public variable
        
    def getName(self): ## Public method
        return self.name

e1 = Employee("Aayush")
print(f"Employee Name: {e1.getName()}")
Employee Name: Aayush

In the case above, both property name and method getName are public attributes.

4.2 Protected Attributes

Protected attributes are similar to public attributes which can be accessed within the class and also available to subclasses. The only difference is the convention, which is to define each protected member with a single underscore “_”.

class Employee:
    def __init__(self, name, project):
        self.name = name ## Public variable
        self._project = project ## Protected variable
        
    def getName(self): ## Public method
        return self.name
    
    def _getProject(self): ## Protected method
        return self._project
    
e1 = Employee("Aayush", "Project Orland")
print(f"Employee Name: {e1.getName()}")
print(f"Project Name: {e1._getProject()}")
Employee Name: Aayush
Project Name: Project Orland

In the case above, both property _project and method _getProject are protected attributes.

4.3 Private Attributes

Private attributes are accessible within the class but not outside of the class. To define a private attribute, prefix the method or property with the double underscore”_“.

class Employee:
    def __init__(self, name, project, salary):
        self.name = name ## Public variable
        self._project = project ## Protected variable
        self.__salary = salary
        
    def getName(self): ## Public method
        return self.name
    
    def _getProject(self): ## Protected method
        return self._project
    
    def __getSalary(self): ## Protected method
        return self.__salary
    
e1 = Employee("Aayush", "Project Orland", "3500")
print(f"Employee Name: {e1.getName()}")
print(f"Project Name: {e1.__getSalary()}") 
Employee Name: Aayush
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [14], in <cell line: 18>()
     16 e1 = Employee("Aayush", "Project Orland", "3500")
     17 print(f"Employee Name: {e1.getName()}")
---> 18 print(f"Project Name: {e1.__getSalary()}")

AttributeError: 'Employee' object has no attribute '__getSalary'

We can see above, __salary property and __getSalary method are both private attributes and when we call them outside of the class they throw an error that the 'Employee' object has no attribute '__getSalary'.

5 Encapsulation

Encapsulation in OOP refers to binding data and the methods to manipulate that data together in a single unit, that is, class. Encapsulation is usually used to hide the state and representation of the object from the outside. A good use of encapsulation is to make all properties private of a class to prevent direct access from outside and use public methods to let the outside world communicate with the class.

class Employee:
    def __init__(self, name, project, salary):
        self.__name = name ## Public variable
        self.__project = project ## Protected variable
        self.__salary = salary
        
    def getName(self): ## Public method
        return self.__name
    
e1 = Employee("Aayush", "Project Orland", "3500")
print(f"Employee Name: {e1.getName()}")
Employee Name: Aayush

Encapsulation has several advantages -

  • Properties of the class can be hidden from the outside world
  • More control over what the outside world can access from the class

A good example of encapsulation would be an access control class based on username and password.

class Auth:
    def __init__(self, userName=None, password=None):
        self.__userName = userName
        self.__password = password
        
    def login(self, userName, password):
        if (self.__userName == userName) and (self.__password == password):
            print (f"Access granted to {userName}")
        else:
            print("Invalid credentials")
            
e1 = Auth("Aayush", "whatever")
e1.login("Aayush", "whatever") ## This will grant access

e1.login("Aayush", "aasdasd") ## This will say invalid creds
e1.__password ## This will raise an error as private properties can't be accessed from outside.
Access granted to Aayush
Invalid credentials
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [4], in <cell line: 16>()
     13 e1.login("Aayush", "whatever") ## This will grant access
     15 e1.login("Aayush", "aasdasd") ## This will say invalid creds
---> 16 e1.__password

AttributeError: 'Auth' object has no attribute '__password'

As we can see above __username and __password are protected properties and can only be used by the class to grand or reject access requests.

6 Inheritance

Inheritance provides a way to create new classes from the existing classes. The new class will inherit all the non-private attributes(properties and methods) from the existing class. The new class can be called a child class and the existing class can be called a parent class.

import math
class Shape:
    def __init__(self, name):
        self.name = name
        
    def getArea(self):
        pass
    
    def printDetails(self):
        print(f"This shape is called {self.name} and area is {self.getArea()}.")
        
class Square(Shape):
    def __init__(self, edge):
        ## calling the constructor from parent class Shape
        Shape.__init__(self, name = "Square")
        self.edge = edge
    
    ## Overiding the getArea function
    def getArea(self):
        return self.edge**2
    
class Circle(Shape):
    def __init__(self, radius):
        ## calling the constructor from parent class Shape
        Shape.__init__(self, name = "Circle")
        self.radius = radius
    ## Overiding the getArea function
    def getArea(self):
        return math.pi * (self.radius**2)
    
obj1 = Square(4)
obj1.printDetails()

obj2 = Circle(3)
obj2.printDetails()
This shape is called Square and area is 16.
This shape is called Circle and area is 28.274333882308138.

We can see above we defined a parent class Shape and then we inherited it to create a Square and Circle child class. While defining the Square and Circle class we overwrote the getArea function pertinent to the class but we used the printDetails function from the parent class to print details about child classes. The more common example in the machine learning world would be to create your own models in Pytorch where we inherit from nn.Module class to create a new model.

6.1 Use of super() Function

super() function comes into play when we implement inheritance. The super() function is used to refer to the parent class without explicitly naming the class. super() function can be used to access parent class properties, calling the parent class, and can be used as initializers. Let’s look at the example above and see how we can modify the Square class to use super() function.

class Shape:
    maxArea = 100
    def __init__(self, name): self.name = name
    def getArea(self): pass
    def printDetails(self): 
        print(f"This shape is called {self.name} and area is {self.getArea()}.")
        

class Square(Shape):
    maxArea = 50
    def __init__(self, edge):
        super().__init__(name = "Square") ## Initializing parent class
        self.edge = edge
    
    def getName(self):
        return super().maxArea
    
    def getArea(self):
        return self.edge**2
    
    def printDetails(self):
        super().printDetails() ## Calling a parent class function
        print(f"Max area from Shape class: {super().maxArea}") ## Accessing parent class property
        print(f"Max area from Square class: {self.maxArea}")

obj1 = Square(4)
obj1.getName()
obj1.printDetails()
This shape is called Square and area is 16.
Max area from Shape class: 100
Max area from Square class: 50

As we can see in the example above we have used -

  • super().__init__ to initialize the parent Shape class
  • super().printDetails() function to use a method from parent class
  • super().maxArea to access a property of a parent class

There are many advantages of inheritance -

  1. Reusability - Inheritance makes the code reusable. Common methods and properties can be stored in a parent class and child classes can inherit these methods.
  2. Modification - Code modification becomes easier if we use inheritance, if we want to make a change in the base class function it will be propagated to the child classes.
  3. Extensibility - We can derive new classes from the old ones by keeping things we need in the derived class.

7 Polymorphism

Polymorphism refers to the same object exhibiting different forms and behaviors. For example consider our shape class which could be a square, rectangle, polygon, etc. Instead of writing multiple functions to get the area of these shapes, we can use a common function like getArea() and implement this function in the derived class.

import math
class Shape:
    def __init__(self, name):
        self.name = name
        
    def getArea(self):
        pass
    
    def printDetails(self):
        print(f"This shape is called {self.name} and area is {self.getArea()}.")
        
class Square(Shape):
    def __init__(self, edge):
        ## calling the constructor from parent class Shape
        Shape.__init__(self, name = "Square")
        self.edge = edge
    
    ## Overiding the getArea function
    def getArea(self):
        return self.edge**2
    
class Circle(Shape):
    def __init__(self, radius):
        ## calling the constructor from parent class Shape
        Shape.__init__(self, name = "Circle")
        self.radius = radius
    ## Overiding the getArea function
    def getArea(self):
        return math.pi * (self.radius**2)
    
obj1 = Square(4)
print(f"Area of this {obj1.name} is {obj1.getArea()}")

obj2 = Circle(3)
print(f"Area of this {obj2.name} is {obj2.getArea()}")
Area of this Square is 16
Area of this Circle is 28.274333882308138

As we can see above there is a pre-defined dummy method called getArea in the Shape class. We override this method in the Square and Circle class. This technique is called method overriding. The advantage of method overriding is that the derived class can write its own specific implementation based on the requirement while using the same function name.

7.1 Abstract base classes

Abstract base classes define a set of methods and properties that a class must implement in order to inherit the parent class. This is a useful technique to enforce that certain functions within the derived class must exist. To define an abstract base class, we use the abc module. The abstract base class inherits from the built-in ABC class and we use the decorator @abstractmethod to declare an abstract method.

from abc import ABC, abstractmethod

class Shape(ABC):
    def __init__(self, name):
        self.name = name
    
    @abstractmethod
    def getArea(self):
        pass
    
    def printDetails(self):
        print(f"This shape is called {self.name} and area is {self.getArea()}.")
        
class Square(Shape):
    def __init__(self, edge):
        ## calling the constructor from parent class Shape
        Shape.__init__(self, name = "Square")
        self.edge = edge
    
obj1 = Square(4)
print(f"Area of this {obj1.name} is {obj1.getArea()}")
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-9-047869063ef7> in <module>
     18         self.edge = edge
     19 
---> 20 obj1 = Square(4)
     21 print(f"Area of this {obj1.name} is {obj1.getArea()}")

TypeError: Can't instantiate abstract class Square with abstract methods getArea

We can see above that we have created a Shape class from the ABC class which has an abstract method getArea. Since our child class Square didn’t have getArea implemented we get an error instantiating this class.

from abc import ABC, abstractmethod

class Shape(ABC):
    def __init__(self, name):
        self.name = name
    
    @abstractmethod
    def getArea(self):
        pass
    
    def printDetails(self):
        print(f"This shape is called {self.name} and area is {self.getArea()}.")
        
class Square(Shape):
    def __init__(self, edge):
        ## calling the constructor from parent class Shape
        Shape.__init__(self, name = "Square")
        self.edge = edge
    
    def getArea(self): return self.edge**2
    
obj1 = Square(4)
print(f"Area of this {obj1.name} is {obj1.getArea()}")
Area of this Square is 16

We can see above, once we implemented the getArea method, the code runs fine.

Abstract base classes serve as a blueprint for derived classes to implement methods that are required to run the function appropriately.

8 Conclusion

In this article, we learned about what is object-oriented programming and key concepts using Python. A good understanding of these concepts will lay a solid foundation for any software professional to write and understand python code better.

I hope you enjoyed reading it. If there is any feedback on the code or just the blog post, feel free to comment below or reach out on LinkedIn.

]]> Programming https://aayushmnit.com/posts/2022-12-20-PythonFundamentals/Python OOPs Fundamentals.html Tue, 20 Dec 2022 08:00:00 GMT Stable diffusion using 🤗 Hugging Face - DiffEdit paper implementation Aayush Agrawal https://aayushmnit.com/posts/2022-11-17-DiffEdit/2022-11-17-DiffEdit.html

An implementation of DIFFEDIT: DIFFUSION-BASED SEMANTIC IMAGE EDITING WITH MASK GUIDANCE using 🤗 hugging face diffusers library.

In this post, I am going to implement a recent paper that came from researchers in Meta AI and Sorbonne Universite named DIFFEDIT. This blog will make more sense to people who are either familiar with the stable diffusion process or are reading after four-part series I made on Stable Diffusion -
1. Part 1 - Stable diffusion using 🤗 Hugging Face - Introduction.
2. Part 2 - Stable diffusion using 🤗 Hugging Face - Looking under the hood.
3. Part 3 - Stable diffusion using 🤗 Hugging Face - Putting everything together
4. Part 4 - Stable diffusion using 🤗 Hugging Face - Variations of Stable Diffusion

Originally, this was the blog post I wanted to write about but realized there is no single place for understanding Stable diffusion with code. Which is the reason I ended up creating the four-part series as a reference or pre-read material to understand this paper.

1 What is DiffEdit?

In simple terms, you can think of DiffEdit approach as a more controlled version of the Image to Image pipeline. DiffEdit takes three inputs-
1. An input image
2. Caption - Describing the input image
3. Target Query - Describe the new image you want to generate

and produce a modified version of the original image based on the query text. This process is particularly good if you want to make a slight tweak to the actual image without completely modifying it.

Fig. 1: Overview of Diff Edit.

As we can see from the image above only the fruits parts of the image were replaced with pears. Pretty amazing results!

The way the authors explain they achieve it is by introducing a mask generation module that determines which part of the image should be edited and then only perform text-based diffusion conditioning on the masked part.

Fig. 2: From the paper DiffEdit. An approach to change an input image by providing caption text and new text.

As we can see from the image above taken from the paper, the authors create a mask from the input image which accurately determines the part of the image where fruits are present and generate a mask (shown in Orange) and then perform masked diffusion to replace fruits with pears. Reading further the authors provide a good visual representation of the whole DiffEdit process.

Fig. 3: Three steps of DiffEdit. Credit - Paper

As I was reading this paper, it seems generating the masking is the most important step and the rest is just textual conditioning using the diffusion process. The conditioning of an image using the mask is a similar idea implemented in Hugging face In-Paint Pipeline. As suggested by the authors, “there are three steps to the DiffEdit process -
Step 1: Add noise to the input image, and denoise it: once conditioned on the query text, and once conditioned on a reference text (or unconditionally). We derive a mask based on the difference in the denoising results.
Step2: we encode the input image with DDIM, to estimate the latents corresponding to the input image
Step3: we perform DDIM decoding conditioned on the text query, using the inferred mask to replace the background with pixel values coming from the encoding process at the corresponding timestep”1

In the next sections, we will start implementing these ideas into actual code.

Let’s start by importing the required libraries and helper functions. All of this was already used and explained in the previous part 2 and part 3 of the stable diffusion series.

Code 
import torch, logging

## disable warnings
logging.disable(logging.WARNING)  

## Imaging  library
from PIL import Image
from torchvision import transforms as tfms


## Basic libraries
from fastdownload import FastDownload
import numpy as np
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
%matplotlib inline
from IPython.display import display
import shutil
import os

## For video display
from IPython.display import HTML
from base64 import b64encode


## Import the CLIP artifacts 
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler

## Helper functions

def load_artifacts():
    '''
    A function to load all diffusion artifacts
    '''
    vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float16).to("cuda")
    unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16).to("cuda")
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16).to("cuda")
    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)    
    return vae, unet, tokenizer, text_encoder, scheduler

def load_image(p):
    '''
    Function to load images from a defined path
    '''
    return Image.open(p).convert('RGB').resize((512,512))

def pil_to_latents(image):
    '''
    Function to convert image to latents
    '''
    init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
    init_image = init_image.to(device="cuda", dtype=torch.float16) 
    init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
    return init_latent_dist

def latents_to_pil(latents):
    '''
    Function to convert latents to images
    '''
    latents = (1 / 0.18215) * latents
    with torch.no_grad():
        image = vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images

def text_enc(prompts, maxlen=None):
    '''
    A function to take a texual promt and convert it into embeddings
    '''
    if maxlen is None: maxlen = tokenizer.model_max_length
    inp = tokenizer(prompts, padding="max_length", max_length=maxlen, truncation=True, return_tensors="pt") 
    return text_encoder(inp.input_ids.to("cuda"))[0].half()

vae, unet, tokenizer, text_encoder, scheduler = load_artifacts()

Let’s also download an image which we will use for the code implementation process.

p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
init_img = load_image(p)
init_img

2 DiffEdit: Purist implementation

Let’s start by implementing the paper as closely as the authors suggested, hence the Purist implementation.

2.1 Mask Creation: First Step of the DiffEdit process

Fig. 4: Step 1 from the DiffEdit paper.

There is a more detailed explanation of Step 1 from the paper, here are the key parts mentioned -
1. Denoise image using different text conditioning, one using reference text and the other using query text, and take differences from the result. The idea is there are more changes in the different parts and not in the background of the image.
2. Repeat this differencing process 10 times
3. Average out these differences and binarize for mask

Note

The third step in mask creation (averaging and binarization) is not explained clearly in the paper and it took me a lot of experiments to get this right.

First, we will try to implement the paper exactly as it’s mentioned. We will modify the prompt_2_img_i2i function for this task to return latents instead of rescaled and decoded de-noised images.

def prompt_2_img_i2i(prompts, init_img, neg_prompts=None, g=7.5, seed=100, strength =0.8, steps=50, dim=512):
    """
    Diffusion process to convert prompt to image
    """
    # Converting textual prompts to embedding
    text = text_enc(prompts) 
    
    # Adding an unconditional prompt , helps in the generation process
    if not neg_prompts: uncond =  text_enc([""], text.shape[1])
    else: uncond =  text_enc(neg_prompt, text.shape[1])
    emb = torch.cat([uncond, text])
    
    # Setting the seed
    if seed: torch.manual_seed(seed)
    
    # Setting number of steps in scheduler
    scheduler.set_timesteps(steps)
    
    # Convert the seed image to latent
    init_latents = pil_to_latents(init_img)
    
    # Figuring initial time step based on strength
    init_timestep = int(steps * strength) 
    timesteps = scheduler.timesteps[-init_timestep]
    timesteps = torch.tensor([timesteps], device="cuda")
    
    # Adding noise to the latents 
    noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
    init_latents = scheduler.add_noise(init_latents, noise, timesteps)
    latents = init_latents
    
    # Computing the timestep to start the diffusion loop
    t_start = max(steps - init_timestep, 0)
    timesteps = scheduler.timesteps[t_start:].to("cuda")
    
    # Iterating through defined steps
    for i,ts in enumerate(tqdm(timesteps)):
        # We need to scale the i/p latents to match the variance
        inp = scheduler.scale_model_input(torch.cat([latents] * 2), ts)
        
        # Predicting noise residual using U-Net
        with torch.no_grad(): u,t = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(2)
            
        # Performing Guidance
        pred = u + g*(t-u)

        # Conditioning  the latents
        #latents = scheduler.step(pred, ts, latents).pred_original_sample
        latents = scheduler.step(pred, ts, latents).prev_sample
    
    # Returning the latent representation to output an array of 4x64x64
    return latents.detach().cpu()

Next, we will make a create_mask function, which will take an initial image, reference prompt, and query prompt with the number of times we need to repeat the steps. In the paper, the author suggests that n=10 and a strength of 0.5 works well in their experimentation. Hence, the default for the function is adjusted to that. create_mask function performs the following steps -
1. Create two denoised latents, one conditioned on reference text and the second on query text, and take a difference of these latents
2. Repeat this step n times
3. Take an average of these differences and standardize
4. Pick a threshold of 0.5 to binarize and create a mask

def create_mask(init_img, rp, qp, n=10, s=0.5):
    ## Initialize a dictionary to save n iterations
    diff = {}
    
    ## Repeating the difference process n times
    for idx in range(n):
        ## Creating denoised sample using reference / original text
        orig_noise = prompt_2_img_i2i(prompts=rp, init_img=init_img, strength=s, seed = 100*idx)[0]
        ## Creating denoised sample using query / target text
        query_noise = prompt_2_img_i2i(prompts=qp, init_img=init_img, strength=s, seed = 100*idx)[0]
        ## Taking the difference 
        diff[idx] = (np.array(orig_noise)-np.array(query_noise))
    
    ## Creating a mask placeholder
    mask = np.zeros_like(diff[0])
    
    ## Taking an average of 10 iterations
    for idx in range(n):
        ## Note np.abs is a key step
        mask += np.abs(diff[idx])  
        
    ## Averaging multiple channels 
    mask = mask.mean(0)
    
    ## Normalizing 
    mask = (mask - mask.mean()) / np.std(mask)
    
    ## Binarizing and returning the mask object
    return (mask > 0).astype("uint8")

mask = create_mask(init_img=init_img, rp=["a horse image"], qp=["a zebra image"], n=10)

Let’s visualize the generated mask over the image.

Code 
plt.imshow(np.array(init_img), cmap='gray') # I would add interpolation='none'
plt.imshow(
    Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
    cmap='cividis', 
    alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
)

Fig. 5: Masking visualization over our horse image.


As we can see above, the mask produced covers the horse portion well which is what we want.

2.2 Masked Diffusion: Step 2 and 3 of DiffEdit paper.

Fig. 6: Step 2 and 3 from the DiffEdit paper.

Steps 2 and 3 need to be implemented in the same loop. Simply put author is saying to condition the latents based on reference text for the non-masked part and on query text for the masked part.
Combine these two parts using this simple formula to create combined latents - 

def prompt_2_img_diffedit(rp, qp, init_img, mask, g=7.5, seed=100, strength =0.7, steps=70, dim=512):
    """
    Diffusion process to convert prompt to image
    """
    # Converting textual prompts to embedding
    rtext = text_enc(rp) 
    qtext = text_enc(qp)
    
    # Adding an unconditional prompt , helps in the generation process
    uncond =  text_enc([""], rtext.shape[1])
    emb = torch.cat([uncond, rtext, qtext])
    
    # Setting the seed
    if seed: torch.manual_seed(seed)
    
    # Setting number of steps in scheduler
    scheduler.set_timesteps(steps)
    
    # Convert the seed image to latent
    init_latents = pil_to_latents(init_img)
    
    # Figuring initial time step based on strength
    init_timestep = int(steps * strength) 
    timesteps = scheduler.timesteps[-init_timestep]
    timesteps = torch.tensor([timesteps], device="cuda")
    
    # Adding noise to the latents 
    noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
    init_latents = scheduler.add_noise(init_latents, noise, timesteps)
    latents = init_latents
    
    # Computing the timestep to start the diffusion loop
    t_start = max(steps - init_timestep, 0)
    timesteps = scheduler.timesteps[t_start:].to("cuda")
    
    # Converting mask to torch tensor
    mask = torch.tensor(mask, dtype=unet.dtype).unsqueeze(0).unsqueeze(0).to("cuda")
    
    # Iterating through defined steps
    for i,ts in enumerate(tqdm(timesteps)):
        # We need to scale the i/p latents to match the variance
        inp = scheduler.scale_model_input(torch.cat([latents] * 3), ts)
        
        # Predicting noise residual using U-Net
        with torch.no_grad(): u, rt, qt = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(3)
            
        # Performing Guidance
        rpred = u + g*(rt-u)
        qpred = u + g*(qt-u)

        # Conditioning  the latents
        rlatents = scheduler.step(rpred, ts, latents).prev_sample
        qlatents = scheduler.step(qpred, ts, latents).prev_sample
        latents = mask*qlatents + (1-mask)*rlatents
    
    # Returning the latent representation to output an array of 4x64x64
    return latents_to_pil(latents)

Let’s visualize the generated image.

output = prompt_2_img_diffedit(
    rp = ["a horse image"], 
    qp=["a zebra image"],
    init_img=init_img, 
    mask = mask, 
    g=7.5, seed=100, strength =0.5, steps=70, dim=512)

## Plotting side by side
fig, axs = plt.subplots(1, 2, figsize=(12, 6))
for c, img in enumerate([init_img, output[0]]): 
    axs[c].imshow(img)
    if c == 0 : axs[c].set_title(f"Initial image ")
    else: axs[c].set_title(f"DiffEdit output")

Fig. 7: DiffEdit output visualization


Let’s create a simple function for the masking and diffusion process.

def diffEdit(init_img, rp , qp, g=7.5, seed=100, strength =0.7, steps=70, dim=512):
    
    ## Step 1: Create mask
    mask = create_mask(init_img=init_img, rp=rp, qp=qp)
    
    ## Step 2 and 3: Diffusion process using mask
    output = prompt_2_img_diffedit(
        rp = rp, 
        qp=qp, 
        init_img=init_img, 
        mask = mask, 
        g=g, 
        seed=seed,
        strength =strength, 
        steps=steps, 
        dim=dim)
    return mask , output

Let’s also create a visualization function for DiffEdit showing the original input image, masked image, and final output image.

def plot_diffEdit(init_img, output, mask):
    ## Plotting side by side
    fig, axs = plt.subplots(1, 3, figsize=(12, 6))
    
    ## Visualizing initial image
    axs[0].imshow(init_img)
    axs[0].set_title(f"Initial image")
    
    ## Visualizing initial image
    axs[2].imshow(output[0])
    axs[2].set_title(f"DiffEdit output")
    
    ## Visualizing the mask 
    axs[1].imshow(np.array(init_img), cmap='gray') 
    axs[1].imshow(
        Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
        cmap='cividis', 
        alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
    )
    axs[1].set_title(f"DiffEdit mask")

Let’s test this function on a few images.

p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
init_img = load_image(p)
mask, output = diffEdit(init_img, rp = ["a horse image"], qp=["a zebra image"])
plot_diffEdit(init_img, output, mask)

Fig. 8: Purist implementation output example


Perfect, let’s try another one.

p = FastDownload().download('https://raw.githubusercontent.com/johnrobinsn/diffusion_experiments/main/images/bowloberries_scaled.jpg')
init_img = load_image(p)
mask, output = diffEdit(init_img, rp = ['Bowl of Strawberries'], qp=['Bowl of Grapes'])
plot_diffEdit(init_img, output, mask)

Fig. 9: Purist implementation output example

3 FastDiffEdit: A faster DiffEdit implementation

Now we have seen the purist implementation, there are some improvements I suggest we can make to the original DiffEdit process in terms of speed and better results. Let’s call these improvements FastDiffEdit.

3.1 Mask Creation: Fast DiffEdit masking process

My biggest issue with the current way of doing masking is that it takes too much time(~50 sec on A4500 GPU). My take is we don’t need to run a full diffusion loop to denoise the image but just use the U-net prediction of the original sample in one shot and increase the repetition to 20 times. In this case, we can improve the computation from 10*25 = 250 steps to 20 steps (12x less loop). Let’s see if this works in practice.

def prompt_2_img_i2i_fast(prompts, init_img, g=7.5, seed=100, strength =0.5, steps=50, dim=512):
    """
    Diffusion process to convert prompt to image
    """
    # Converting textual prompts to embedding
    text = text_enc(prompts) 
    
    # Adding an unconditional prompt , helps in the generation process
    uncond =  text_enc([""], text.shape[1])
    emb = torch.cat([uncond, text])
    
    # Setting the seed
    if seed: torch.manual_seed(seed)
    
    # Setting number of steps in scheduler
    scheduler.set_timesteps(steps)
    
    # Convert the seed image to latent
    init_latents = pil_to_latents(init_img)
    
    # Figuring initial time step based on strength
    init_timestep = int(steps * strength) 
    timesteps = scheduler.timesteps[-init_timestep]
    timesteps = torch.tensor([timesteps], device="cuda")
    
    # Adding noise to the latents 
    noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
    init_latents = scheduler.add_noise(init_latents, noise, timesteps)
    latents = init_latents
    
    # We need to scale the i/p latents to match the variance
    inp = scheduler.scale_model_input(torch.cat([latents] * 2), timesteps)
    # Predicting noise residual using U-Net
    with torch.no_grad(): u,t = unet(inp, timesteps, encoder_hidden_states=emb).sample.chunk(2)
         
    # Performing Guidance
    pred = u + g*(t-u)

    # Zero shot prediction
    latents = scheduler.step(pred, timesteps, latents).pred_original_sample
    
    # Returning the latent representation to output an array of 4x64x64
    return latents.detach().cpu()

Let’s create a new masking function that can take our prompt_2_img_i2i_fast function.

def create_mask_fast(init_img, rp, qp, n=20, s=0.5):
    ## Initialize a dictionary to save n iterations
    diff = {}
    
    ## Repeating the difference process n times
    for idx in range(n):
        ## Creating denoised sample using reference / original text
        orig_noise = prompt_2_img_i2i_fast(prompts=rp, init_img=init_img, strength=s, seed = 100*idx)[0]
        ## Creating denoised sample using query / target text
        query_noise = prompt_2_img_i2i_fast(prompts=qp, init_img=init_img, strength=s, seed = 100*idx)[0]
        ## Taking the difference 
        diff[idx] = (np.array(orig_noise)-np.array(query_noise))
    
    ## Creating a mask placeholder
    mask = np.zeros_like(diff[0])
    
    ## Taking an average of 10 iterations
    for idx in range(n):
        ## Note np.abs is a key step
        mask += np.abs(diff[idx])  
        
    ## Averaging multiple channels 
    mask = mask.mean(0)
    
    ## Normalizing 
    mask = (mask - mask.mean()) / np.std(mask)
    
    ## Binarizing and returning the mask object
    return (mask > 0).astype("uint8")

Let’s see if this new masking function produces a good mask.

p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
init_img = load_image(p)
mask = create_mask_fast(init_img=init_img, rp=["a horse image"], qp=["a zebra image"], n=20)
plt.imshow(np.array(init_img), cmap='gray') # I would add interpolation='none'
plt.imshow(
    Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
    cmap='cividis', 
    alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
)

Fig. 10: FastDiffEdit masking visualization over our horse image.


As we can see above the masking is improved and compute time has reduced from ~50 seconds to ~10 secs on my machine(5x improvement!).

Let’s improve our masking by adding a cv2 trick. This will just smooth out the masking a little bit more.

import cv2
def improve_mask(mask):
    mask  = cv2.GaussianBlur(mask*255,(3,3),1) > 0
    return mask.astype('uint8')
mask = improve_mask(mask)
plt.imshow(np.array(init_img), cmap='gray') # I would add interpolation='none'
plt.imshow(
    Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
    cmap='cividis', 
    alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
)

Fig. 11: Improved FastDiffEdit masking visualization over our horse image with cv2 Gaussian blur trick.


As we can see above the masking has become a bit more smooth and covers more area.

3.2 Masked Diffusion: Replace with 🤗 inpaint pipeline

So, instead of using our function to perform the masked diffusion, there is a special pipeline in 🤗 diffusers library called inpaint pipeline. Which takes the query prompt, initial image, and generated mask to generate the output image. Let’s start by loading in the inpaint pipeline.

from diffusers import StableDiffusionInpaintPipeline
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    revision="fp16",
    torch_dtype=torch.float16,
).to("cuda")

Let’s use the inpaint pipeline with our generated mask and image.

pipe(
    prompt=["a zebra image"], 
    image=init_img, 
    mask_image=Image.fromarray(mask*255).resize((512,512)), 
    generator=torch.Generator("cuda").manual_seed(100),
    num_inference_steps = 20
).images[0]
image

Fig. 12: In-paint pipeline output.


As we can see above, inpaint pipeline creates a more realistic zebra image. Let’s create a simple function for the masking and diffusion process.

def fastDiffEdit(init_img, rp , qp, g=7.5, seed=100, strength =0.7, steps=20, dim=512):
    
    ## Step 1: Create mask
    mask = create_mask_fast(init_img=init_img, rp=rp, qp=qp, n=20)
    
    ## Improve masking using CV trick
    mask = improve_mask(mask)
    
    ## Step 2 and 3: Diffusion process using mask
    output = pipe(
        prompt=qp, 
        image=init_img, 
        mask_image=Image.fromarray(mask*255).resize((512,512)), 
        generator=torch.Generator("cuda").manual_seed(100),
        num_inference_steps = steps
    ).images
    return mask , output

Let’s test this function on a few images.

p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
init_img = load_image(p)
mask, output = fastDiffEdit(init_img, rp = ["a horse image"], qp=["a zebra image"])
plot_diffEdit(init_img, output, mask)

Fig. 13: FastDiffEdit output example

Perfect, let’s try another one.

p = FastDownload().download('https://raw.githubusercontent.com/johnrobinsn/diffusion_experiments/main/images/bowloberries_scaled.jpg')
init_img = load_image(p)
mask, output = fastDiffEdit(init_img, rp = ['Bowl of Strawberries'], qp=['Bowl of Grapes'])
plot_diffEdit(init_img, output, mask)

Fig. 14: FastDiffEdit output example

4 Conclusion

In this post, we implemented the DiffEdit paper as the author mentioned and then we proposed improvements to the method to create FastDiffEdit which speeds up computation times up to 5 times.

I hope you enjoyed reading it, and feel free to use my code and try it out for generating your images. Also, if there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn or email me at aayushmnit@gmail.com.

]]> Stable Diffusion Research https://aayushmnit.com/posts/2022-11-17-DiffEdit/2022-11-17-DiffEdit.html Thu, 17 Nov 2022 08:00:00 GMT Stable diffusion using 🤗 Hugging Face - Variations of Stable Diffusion Aayush Agrawal https://aayushmnit.com/posts/2022-11-10-StableDiffusionP4/2022-11-10-StableDiffusionP4.html

An introduction to negative prompting and image to image stable diffusion pipeline using 🤗 hugging face diffusers library.

This is my fourth post of the Stable diffusion series, if you haven’t checked out the previous ones, you can read it here -
1. Part 1 - Stable diffusion using 🤗 Hugging Face - Introduction.
2. Part 2 - Stable diffusion using 🤗 Hugging Face - Looking under the hood.
3. Part 3 - Stable diffusion using 🤗 Hugging Face - Putting everything together

In previous posts, I went over all the key components of Stable Diffusion and how to get a prompt to image pipeline working. In this post, I will show how to edit the prompt to image function to add additional functionality to our Stable diffusion pipeline i.e., Negative prompting and Image to Image pipeline. Hopefully, this will provide enough motivation to play around with this function and conduct your research.

Fig. 1: A Stable diffusion generated image using prompt -
“A road diverging in two different direction”

1 Variation 1: Negative Prompt

1.1 What is negative prompting?

A negative prompt is an additional capability we can add to our model to tell the stable diffusion model what we don’t want to see in the generated image. This feature is popular to remove anything a user doesn’t want to see from the original generated image.

Fig. 2: Negative prompt example

1.2 Understanding negative prompting through code

Let’s start by importing the required libraries and helper functions. All of this was already used and explained in the previous part 2 and part 3 of the series.

Code 
import torch, logging

## disable warnings
logging.disable(logging.WARNING)  

## Imaging  library
from PIL import Image
from torchvision import transforms as tfms


## Basic libraries
from fastdownload import FastDownload
import numpy as np
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
%matplotlib inline
from IPython.display import display
import shutil
import os

## For video display
from IPython.display import HTML
from base64 import b64encode


## Import the CLIP artifacts 
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler

## Initiating tokenizer and encoder.
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16).to("cuda")

## Initiating the VAE
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float16).to("cuda")

## Initializing a scheduler and Setting number of sampling steps
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
scheduler.set_timesteps(50)

## Initializing the U-Net model
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16).to("cuda")

## Helper functions
def load_image(p):
    '''
    Function to load images from a defined path
    '''
    return Image.open(p).convert('RGB').resize((512,512))

def pil_to_latents(image):
    '''
    Function to convert image to latents
    '''
    init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
    init_image = init_image.to(device="cuda", dtype=torch.float16) 
    init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
    return init_latent_dist

def latents_to_pil(latents):
    '''
    Function to convert latents to images
    '''
    latents = (1 / 0.18215) * latents
    with torch.no_grad():
        image = vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images

def text_enc(prompts, maxlen=None):
    '''
    A function to take a texual promt and convert it into embeddings
    '''
    if maxlen is None: maxlen = tokenizer.model_max_length
    inp = tokenizer(prompts, padding="max_length", max_length=maxlen, truncation=True, return_tensors="pt") 
    return text_encoder(inp.input_ids.to("cuda"))[0].half()

Now we are going to change the prompt_2_img function from part 3 by passing an additional function neg_prompts. The way negative prompt works is by using user-specified text instead of an empty string for unconditional embedding(uncond) when doing sampling.

Fig. 3: Negative prompt code change

So let’s make this change and update our prompt_2_img function.

def prompt_2_img(prompts, neg_prompts=None, g=7.5, seed=100, steps=70, dim=512, save_int=False):
    """
    Diffusion process to convert prompt to image
    """
    
    # Defining batch size
    bs = len(prompts) 
    
    # Converting textual prompts to embedding
    text = text_enc(prompts) 
    
    # Adding an unconditional prompt , helps in the generation process
    if not neg_prompts: uncond =  text_enc([""] * bs, text.shape[1])
    else: uncond =  text_enc(neg_prompts, text.shape[1])
    emb = torch.cat([uncond, text])
    
    # Setting the seed
    if seed: torch.manual_seed(seed)
    
    # Initiating random noise
    latents = torch.randn((bs, unet.in_channels, dim//8, dim//8))
    
    # Setting number of steps in scheduler
    scheduler.set_timesteps(steps)
    
    # Adding noise to the latents 
    latents = latents.to("cuda").half() * scheduler.init_noise_sigma
    
    # Iterating through defined steps
    for i,ts in enumerate(tqdm(scheduler.timesteps)):
        # We need to scale the i/p latents to match the variance
        inp = scheduler.scale_model_input(torch.cat([latents] * 2), ts)
        
        # Predicting noise residual using U-Net
        with torch.no_grad(): u,t = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(2)
            
        # Performing Guidance
        pred = u + g*(t-u)
        
        # Conditioning  the latents
        latents = scheduler.step(pred, ts, latents).prev_sample
        
        # Saving intermediate images
        if save_int: 
            if not os.path.exists(f'./steps'): os.mkdir(f'./steps')
            latents_to_pil(latents)[0].save(f'steps/{i:04}.jpeg')
            
    # Returning the latent representation to output an image of 3x512x512
    return latents_to_pil(latents)

Let’s see if the function works as intended.

## Image without neg prompt
images = [None, None]
images[0] = prompt_2_img(prompts = ["A dog wearing a white hat"], neg_prompts=[""],steps=50, save_int=False)[0]
images[1] = prompt_2_img(prompts = ["A dog wearing a white hat"], neg_prompts=["White hat"],steps=50, save_int=False)[0]
    
## Plotting side by side
fig, axs = plt.subplots(1, 2, figsize=(12, 6))
for c, img in enumerate(images): 
    axs[c].imshow(img)
    if c == 0 : axs[c].set_title(f"A dog wearing a white hat")
    else: axs[c].set_title(f"Neg prompt - white hat")

Fig. 4: Visualization of negative prompting. Left SD generated with prompt “A dog wearing a white hat” and on right the same caption with negative prompt of “White hat”


As we can see it can be a really handy feature to fine-tune the image to your liking. You can also use it to generate a pretty realistic face by being really descriptive as this Reddit post. Let’s try it -

prompt = ['Close-up photography of the face of a 30 years old man with brown eyes, (by Alyssa Monks:1.1), by Joseph Lorusso, by Lilia Alvarado, beautiful lighting, sharp focus, 8k, high res, (pores:0.1), (sweaty:0.8), Masterpiece, Nikon Z9, Award - winning photograph']
neg_prompt = ['lowres, signs, memes, labels, text, food, text, error, mutant, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, made by children, caricature, ugly, boring, sketch, lacklustre, repetitive, cropped, (long neck), facebook, youtube, body horror, out of frame, mutilated, tiled, frame, border, porcelain skin, doll like, doll']
images = prompt_2_img(prompts = prompt, neg_prompts=neg_prompt, steps=50, save_int=False)
images[0]

Fig. 5: An image generated using negative prompting.


Pretty neat! I hope this gives you some ideas on how to get going with your own variations of stable diffusion. Now let’s look at another variation of Stable diffusion.

2 Variation 2: Image to Image pipeline

2.1 What is an image to image pipeline?

As seen above, prompt_2_img functions start generating an image from random gaussian noise, but what if we feed an initial seed image to guide the diffusion process? This is exactly how the image to image pipeline works. Instead of purely relying on text conditioning for the output image, we can use an initial seed image mix it with some noise (which can be guided by a strength parameter), and then run the diffusion loop.

Fig. 6: Image to image pipeline example.

2.2 Understanding image to image prompting through code

Now we are going to change the prompt_2_img function defined above. We will introduce two more parameters to our prompt_2_img_i2i function -
1. init_img: Which is going to be the Image object containing the seed image
2. strength: This parameter will take a value between 0 and 1. The higher the value less the final image is going to look similar to the seed image.

def prompt_2_img_i2i(prompts, init_img, neg_prompts=None, g=7.5, seed=100, strength =0.8, steps=50, dim=512, save_int=False):
    """
    Diffusion process to convert prompt to image
    """
    # Converting textual prompts to embedding
    text = text_enc(prompt) 
    
    # Adding an unconditional prompt , helps in the generation process
    if not neg_prompts: uncond =  text_enc([""], text.shape[1])
    else: uncond =  text_enc(neg_prompt, text.shape[1])
    emb = torch.cat([uncond, text])
    
    # Setting the seed
    if seed: torch.manual_seed(seed)
    
    # Setting number of steps in scheduler
    scheduler.set_timesteps(steps)
    
    # Convert the seed image to latent
    init_latents = pil_to_latents(init_img)
    
    # Figuring initial time step based on strength
    init_timestep = int(steps * strength) 
    timesteps = scheduler.timesteps[-init_timestep]
    timesteps = torch.tensor([timesteps], device="cuda")
    
    # Adding noise to the latents 
    noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
    init_latents = scheduler.add_noise(init_latents, noise, timesteps)
    latents = init_latents
    
    # Computing the timestep to start the diffusion loop
    t_start = max(steps - init_timestep, 0)
    timesteps = scheduler.timesteps[t_start:].to("cuda")
    
    # Iterating through defined steps
    for i,ts in enumerate(tqdm(timesteps)):
        # We need to scale the i/p latents to match the variance
        inp = scheduler.scale_model_input(torch.cat([latents] * 2), ts)
        
        # Predicting noise residual using U-Net
        with torch.no_grad(): u,t = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(2)
            
        # Performing Guidance
        pred = u + g*(t-u)
        
        # Conditioning  the latents
        latents = scheduler.step(pred, ts, latents).prev_sample
        
        # Saving intermediate images
        if save_int: 
            if not os.path.exists(f'./steps'):
                os.mkdir(f'./steps')
            latents_to_pil(latents)[0].save(f'steps/{i:04}.jpeg')
            
    # Returning the latent representation to output an image of 3x512x512
    return latents_to_pil(latents)

Instead of using random noise, you will notice we use the strength parameter to figure out how much noise to add and also the number of steps to run the diffusion loop for. The amount of noise is calculated by multiplying strength(default = 0.8) with steps (default = 50) which is the 10th (50 - 50 * 0.8) step and running the diffusion loop for 40(50*0.8) remaining steps. Let’s load an initial image and pass it through the prompt_2_img_i2i function.

p = FastDownload().download('https://s3.amazonaws.com/moonup/production/uploads/1664665907257-noauth.png')
image = Image.open(p).convert('RGB').resize((512,512))
prompt = ["Wolf howling at the moon, photorealistic 4K"]
images = prompt_2_img_i2i(prompts = prompt, init_img = image)
/home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages/diffusers/schedulers/scheduling_lms_discrete.py:146: IntegrationWarning: The maximum number of subdivisions (50) has been achieved.
  If increasing the limit yields no improvement it is advised to analyze 
  the integrand in order to determine the difficulties.  If the position of a 
  local difficulty can be determined (singularity, discontinuity) one will 
  probably gain from splitting up the interval and calling the integrator 
  on the subranges.  Perhaps a special-purpose integrator should be used.
  integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0]
Code 
## Plotting side by side
fig, axs = plt.subplots(1, 2, figsize=(12, 6))
for c, img in enumerate([image, images[0]]): 
    axs[c].imshow(img)
    if c == 0 : axs[c].set_title(f"Initial image")
    else: axs[c].set_title(f"Image 2 Image output")

Fig. 7: Visualization of image to image pipeline. Left is initial image passed in img2img pipeline and right is the output of the img2img pipeline.


We can see our prompt_2_img_i2i function creates a pretty epic image from the initial sketch provided.

3 Conclusion

I hope this gives a good overview of how to tweak the prompt_2_img function to add additional capabilities to your stable diffusion loop. The understanding of this lower-level function is useful for trying your own idea to improve stable diffusion or implement new papers which I might cover in my next post.

I hope you enjoyed reading it, and feel free to use my code and try it out for generating your images. Also, if there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn or email me at aayushmnit@gmail.com.

4 References

]]> Stable Diffusion https://aayushmnit.com/posts/2022-11-10-StableDiffusionP4/2022-11-10-StableDiffusionP4.html Fri, 11 Nov 2022 08:00:00 GMT Stable diffusion using 🤗 Hugging Face - Putting everything together Aayush Agrawal https://aayushmnit.com/posts/2022-11-07-StableDiffusionP3/2022-11-07-StableDiffusionP3.html

An introduction to the diffusion process using 🤗 hugging face diffusers library.

This is my third post of the Stable diffusion series, if you haven’t checked out the previous ones, you can read it here -
1. Part 1 - Stable diffusion using 🤗 Hugging Face - Introduction.
2. Part 2 - Stable diffusion using 🤗 Hugging Face - Looking under the hood.

In previous posts, I went over showing how to install 🤗 diffuser library to start generating your own AI images and key components of the stable diffusion pipeline i.e., CLIP text encoder, VAE, and U-Net. In this post, we will try to put these key components together and do a walk-through of the diffusion process which generates the image.

Fig. 1: Building Blocks Vector

1 Overview - The Diffusion Process

The stable diffusion model takes the textual input and a seed. The textual input is then passed through the CLIP model to generate textual embedding of size 77x768 and the seed is used to generate Gaussian noise of size 4x64x64 which becomes the first latent image representation.

Note

You will notice that there is an additional dimension mentioned (1x) in the image like 1x77x768 for text embedding, that is because it represents the batch size of 1.

Fig. 2: The diffusion process.

Next, the U-Net iteratively denoises the random latent image representations while conditioning on the text embeddings. The output of the U-Net is predicted noise residual, which is then used to compute conditioned latents via a scheduler algorithm. This process of denoising and text conditioning is repeated N times (We will use 50) to retrieve a better latent image representation. Once this process is complete, the latent image representation (4x64x64) is decoded by the VAE decoder to retrieve the final output image (3x512x512).

Note

This iterative denoising is an important step for getting a good output image. Typical steps are in the range of 30-80. However, there are recent papers that claim to reduce it to 4-5 steps by using distillation techniques.

2 Understanding the diffusion process through code

Let’s start by importing the required libraries and helper functions. All of this was already used and explained in the previous part 2 of the series.

import torch, logging

## disable warnings
logging.disable(logging.WARNING)  

## Imaging  library
from PIL import Image
from torchvision import transforms as tfms

## Basic libraries
import numpy as np
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
%matplotlib inline
from IPython.display import display
import shutil
import os

## For video display
from IPython.display import HTML
from base64 import b64encode


## Import the CLIP artifacts 
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler

## Initiating tokenizer and encoder.
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16).to("cuda")

## Initiating the VAE
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float16).to("cuda")

## Initializing a scheduler and Setting number of sampling steps
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
scheduler.set_timesteps(50)

## Initializing the U-Net model
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16).to("cuda")

## Helper functions
def load_image(p):
    '''
    Function to load images from a defined path
    '''
    return Image.open(p).convert('RGB').resize((512,512))

def pil_to_latents(image):
    '''
    Function to convert image to latents
    '''
    init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
    init_image = init_image.to(device="cuda", dtype=torch.float16) 
    init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
    return init_latent_dist

def latents_to_pil(latents):
    '''
    Function to convert latents to images
    '''
    latents = (1 / 0.18215) * latents
    with torch.no_grad():
        image = vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images

def text_enc(prompts, maxlen=None):
    '''
    A function to take a texual promt and convert it into embeddings
    '''
    if maxlen is None: maxlen = tokenizer.model_max_length
    inp = tokenizer(prompts, padding="max_length", max_length=maxlen, truncation=True, return_tensors="pt") 
    return text_encoder(inp.input_ids.to("cuda"))[0].half()

The code below is a stripped-down version of what is present in the StableDiffusionPipeline.from_pretrained function to show the important parts of the diffusion process.

def prompt_2_img(prompts, g=7.5, seed=100, steps=70, dim=512, save_int=False):
    """
    Diffusion process to convert prompt to image
    """
    
    # Defining batch size
    bs = len(prompts) 
    
    # Converting textual prompts to embedding
    text = text_enc(prompts) 
    
    # Adding an unconditional prompt , helps in the generation process
    uncond =  text_enc([""] * bs, text.shape[1])
    emb = torch.cat([uncond, text])
    
    # Setting the seed
    if seed: torch.manual_seed(seed)
    
    # Initiating random noise
    latents = torch.randn((bs, unet.in_channels, dim//8, dim//8))
    
    # Setting number of steps in scheduler
    scheduler.set_timesteps(steps)
    
    # Adding noise to the latents 
    latents = latents.to("cuda").half() * scheduler.init_noise_sigma
    
    # Iterating through defined steps
    for i,ts in enumerate(tqdm(scheduler.timesteps)):
        # We need to scale the i/p latents to match the variance
        inp = scheduler.scale_model_input(torch.cat([latents] * 2), ts)
        
        # Predicting noise residual using U-Net
        with torch.no_grad(): u,t = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(2)
            
        # Performing Guidance
        pred = u + g*(t-u)
        
        # Conditioning  the latents
        latents = scheduler.step(pred, ts, latents).prev_sample
        
        # Saving intermediate images
        if save_int: 
            if not os.path.exists(f'./steps'):
                os.mkdir(f'./steps')
            latents_to_pil(latents)[0].save(f'steps/{i:04}.jpeg')
            
    # Returning the latent representation to output an image of 3x512x512
    return latents_to_pil(latents)

Let’s see if the function works as intended.

images = prompt_2_img(["A dog wearing a hat", "a photograph of an astronaut riding a horse"], save_int=False)
for img in images:display(img)

Looks like it is working! So let’s take a deeper dive at the hyper-parameters of the function.
1. prompt - this is the textual prompt we pass through to generate an image. Similar to the pipe(prompt) function we saw in part 1
2. g or guidance scale - It’s a value that determines how close the image should be to the textual prompt. This is related to a technique called Classifier free guidance which improves the quality of the images generated. The higher the value of the guidance scale, more close it will be to the textual prompt
3. seed - This sets the seed from which the initial Gaussian noisy latents are generated
4. steps - Number of de-noising steps taken for generating the final latents.
5. dim - dimension of the image, for simplicity we are currently generating square images, so only one value is needed
6. save_int - This is optional, a boolean flag, if we want to save intermediate latent images, helps in visualization.

Let’s visualize this process of generation from noise to the final image.

Code 
## Creating image through prompt_2_img modified function
images = prompt_2_img(["A dog wearing a hat"], save_int=True)

## Converting intermediate images to video
!ffmpeg -v 1 -y -f image2 -framerate 20 -i steps/%04d.jpeg -c:v libx264 -preset slow -qp 18 -pix_fmt yuv420p out.mp4

## Deleting intermediate images
shutil.rmtree(f'./steps/')

## Displaying video output
mp4 = open('out.mp4','rb').read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML("""
<video width=600 controls>
      <source src="%s" type="video/mp4">
</video>
""" % data_url)
Fig 3: The de-noisation steps visualization.

3 Conclusion

I hope this gives a good overview and breaks the code to the bare minimum so that we can understand each component. Now that we have the minimum code implemented, in the next post we will see make some tweaks to the mk_img function to add additional functionality i.e., img2img pipeline and negative prompt.

I hope you enjoyed reading it, and feel free to use my code and try it out for generating your images. Also, if there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn or email me at aayushmnit@gmail.com.

4 References

]]> Stable Diffusion https://aayushmnit.com/posts/2022-11-07-StableDiffusionP3/2022-11-07-StableDiffusionP3.html Mon, 07 Nov 2022 08:00:00 GMT Stable diffusion using 🤗 Hugging Face - Looking under the hood Aayush Agrawal https://aayushmnit.com/posts/2022-11-05-StableDiffusionP2/2022-11-05-StableDiffusionP2.html

An introduction into what goes on in the pipe function of 🤗 hugging face diffusers library StableDiffusionPipeline function.

This is my second post of the Stable diffusion series, if you haven’t checked out the first one, you can read it here -
1. Part 1 - Introduction to Stable diffusion using 🤗 Hugging Face.

In this post, we will understand the basic components of a stable diffusion pipeline and their purpose. Later we will reconstruct StableDiffusionPipeline.from_pretrained function using these components. Let’s get started -

Fig. 1: This image was generated by 🤗 Stable diffusion model using the “a scientist looking under the hood of a car realistic 4k image” prompt.

1 Introduction

Diffusion models as seen in the previous post can generate high-quality images. Stable diffusion models are a special kind of diffusion model called the Latent Diffusion model. They have first proposed in this paper High-Resolution Image Synthesis with Latent Diffusion Models. The original Diffusion model tends to consume a lot more memory, so latent diffusion models were created which can do the diffusion process in lower dimension space called Latent Space. On a high level, diffusion models are machine learning models that are trained to denoise random Gaussian noise step by step, to get the result i.e., image. In latent diffusion, the model is trained to do this same process in a lower dimension.

There are three main components in latent diffusion -

  1. A text encoder, in this case, a CLIP Text encoder
  2. An autoencoder, in this case, a Variational Auto Encoder also referred to as VAE
  3. A U-Net

Let’s dive into each of these components and understand their use in the diffusion process. The way I will be attempting to explain these components is by talking about them in the following three stages -

  1. The Basics: What goes in the component and what comes out of the component - This is an important, and key part of the top-down learning approach of understanding “the whole game”
  2. Deeper explanation using 🤗 code. - This part will provide more understanding of what the model produces using the code
  3. What’s their role in the Stable diffusion pipeline - This will build your intuition around how this component fits in the Stable diffusion process. This will help your intuition on the diffusion process

2 CLIP Text Encoder

2.1 Basics - What goes in and out of the component?

CLIP(Contrastive Language–Image Pre-training) text encoder takes the text as an input and generates text embeddings that are close in latent space as it may be if you would have encoded an image through a CLIP model.

Fig. 2: CLIP text encoder

2.2 Deeper explanation using 🤗 code

Any machine learning model doesn’t understand text data. For any model to understand text data, we need to convert this text into numbers that hold the meaning of the text, referred to as embeddings. The process of converting a text to a number can be broken down into two parts -
1. Tokenizer - Breaking down each word into sub-words and then using a lookup table to convert them into a number
2. Token_To_Embedding Encoder - Converting those numerical sub-words into a representation that contains the representation of that text

Let’s look at it through code. We will start by importing the relevant artifacts.

import torch, logging

## disable warnings
logging.disable(logging.WARNING)  

## Import the CLIP artifacts 
from transformers import CLIPTextModel, CLIPTokenizer

## Initiating tokenizer and encoder.
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16).to("cuda")

Let’s initialize a prompt and tokenize it.

prompt = ["a dog wearing hat"]
tok =tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") 
print(tok.input_ids.shape)
tok
torch.Size([1, 77])
{'input_ids': tensor([[49406,   320,  1929,  3309,  3801, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0]])}

A tokenizer returns two objects in the form of a dictionary -
1. input_ids - A tensor of size 1x77 as one prompt was passed and padded to 77 max length. 49406 is a start token, 320 is a token given to the word “a”, 1929 to the word dog, 3309 to the word wearing, 3801 to the word hat, and 49407 is the end of text token repeated till the pad length of 77.
2. attention_mask - 1 representing an embedded value and 0 representing padding.

for token in list(tok.input_ids[0,:7]): print(f"{token}:{tokenizer.convert_ids_to_tokens(int(token))}")
49406:<|startoftext|>
320:a</w>
1929:dog</w>
3309:wearing</w>
3801:hat</w>
49407:<|endoftext|>
49407:<|endoftext|>

So, let’s look at the Token_To_Embedding Encoder which takes the input_ids generated by the tokenizer and converts them into embeddings -

emb = text_encoder(tok.input_ids.to("cuda"))[0].half()
print(f"Shape of embedding : {emb.shape}")
emb
Shape of embedding : torch.Size([1, 77, 768])
tensor([[[-0.3887,  0.0229, -0.0522,  ..., -0.4902, -0.3066,  0.0673],
         [ 0.0292, -1.3242,  0.3074,  ..., -0.5264,  0.9766,  0.6655],
         [-1.5928,  0.5063,  1.0791,  ..., -1.5283, -0.8438,  0.1597],
         ...,
         [-1.4688,  0.3113,  1.1670,  ...,  0.3755,  0.5366, -1.5049],
         [-1.4697,  0.3000,  1.1777,  ...,  0.3774,  0.5420, -1.5000],
         [-1.4395,  0.3137,  1.1982,  ...,  0.3535,  0.5400, -1.5488]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)

As we can see above, each tokenized input of size 1x77 has now been translated to 1x77x768 shape embedding. So, each word got represented in a 768-dimensional space.

2.3 What’s their role in the Stable diffusion pipeline

Stable diffusion only uses a CLIP trained encoder for the conversion of text to embeddings. This becomes one of the inputs to the U-net. On a high level, CLIP uses an image encoder and text encoder to create embeddings that are similar in latent space. This similarity is more precisely defined as a Contrastive objective. For more information on how CLIP is trained, please refer to this Open AI blog.

Fig. 3: CLIP pre-trains an image encoder and a text encoder to predict which images were paired with which texts in our dataset. Credit - OpenAI

3 VAE - Variational Auto Encoder

3.1 Basics - What goes in and out of the component?

An autoencoder contains two parts -
1. Encoder takes an image as input and converts it into a low dimensional latent representation
2. Decoder takes the latent representation and converts it back into an image

Fig. 4: A Variational autoencoder. Original bird pic credit.

As we can see above, the Encoder acts like a compressor that squishes the image into lower dimensions and the decoder recreates the original image back from the compressed version.

Note

Encoder-Decoder compression-decompression is not lossless.

3.2 Deeper explanation using 🤗 code

Let’s start looking at VAE through code. We will start by importing the required libraries and defining some helper functions.

Code 
## To import an image from a URL
from fastdownload import FastDownload

## Imaging  library
from PIL import Image
from torchvision import transforms as tfms

## Basic libraries
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

## Loading a VAE model
from diffusers import AutoencoderKL
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float16).to("cuda")

def load_image(p):
    '''
    Function to load images from a defined path
    '''
    return Image.open(p).convert('RGB').resize((512,512))

def pil_to_latents(image):
    '''
    Function to convert image to latents
    '''
    init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
    init_image = init_image.to(device="cuda", dtype=torch.float16) 
    init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
    return init_latent_dist

def latents_to_pil(latents):
    '''
    Function to convert latents to images
    '''
    latents = (1 / 0.18215) * latents
    with torch.no_grad():
        image = vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images

Let’s download an image from the internet.

p = FastDownload().download('https://lafeber.com/pet-birds/wp-content/uploads/2018/06/Scarlet-Macaw-2.jpg')
img = load_image(p)
print(f"Dimension of this image: {np.array(img).shape}")
img
Dimension of this image: (512, 512, 3)

Fig. 5: Original bird pic credit.

Now let’s compress this image by using the VAE encoder, we will be using the pil_to_latents helper function.

latent_img = pil_to_latents(img)
print(f"Dimension of this latent representation: {latent_img.shape}")
Dimension of this latent representation: torch.Size([1, 4, 64, 64])

As we can see how the VAE compressed a 3 x 512 x 512 dimension image into a 4 x 64 x 64 image. That’s a compression ratio of 48x! Let’s visualize these four channels of latent representations.

fig, axs = plt.subplots(1, 4, figsize=(16, 4))
for c in range(4):
    axs[c].imshow(latent_img[0][c].detach().cpu(), cmap='Greys')

Fig. 6: Visualization of latent representation from VAE encoder.

This latent representation in theory should capture a lot of information about the original image. Let’s use the decoder on this representation to see what we get back. For this, we will use the latents_to_pil helper function.

decoded_img = latents_to_pil(latent_img)
decoded_img[0]

Fig. 7: Visualization of decoded latent representation from VAE decoder.

As we can see from the figure above VAE decoder was able to recover the original image from a 48x compressed latent representation. That’s impressive!

Note

If you look closely at the decoded image, it’s not the same as the original image, notice the difference around the eyes. That’s why VAE encoder/decoder is not a lossless compression.

3.3 What’s their role in the Stable diffusion pipeline

Stable diffusion can be done without the VAE component but the reason we use VAE is to reduce the computational time to generate High-resolution images. The latent diffusion models can perform diffusion in this latent space produced by the VAE encoder and once we have our desired latent outputs produced by the diffusion process, we can convert them back to the high-resolution image by using the VAE decoder. To get a better intuitive understanding of Variation Autoencoders and how they are trained, read this blog by Irhum Shafkat.

4 U-Net

4.1 Basics - What goes in and out of the component?

The U-Net model takes two inputs -
1. Noisy latent or Noise- Noisy latents are latents produced by a VAE encoder (in case an initial image is provided) with added noise or it can take pure noise input in case we want to create a random new image based solely on a textual description
2. Text embeddings - CLIP-based embedding generated by input textual prompts

The output of the U-Net model is the predicted noise residual which the input noisy latent contains. In other words, it predicts the noise which is subtracted from the noisy latents to return the original de-noised latents.

Fig. 8: A U-Net representation.

4.2 Deeper explanation using 🤗 code

Let’s start looking at U-Net through code. We will start by importing the required libraries and initiating our U-Net model.

from diffusers import UNet2DConditionModel, LMSDiscreteScheduler

## Initializing a scheduler
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
## Setting number of sampling steps
scheduler.set_timesteps(51)

## Initializing the U-Net model
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16).to("cuda")

As you may have noticed from code above, we not only imported unet but also a scheduler. The purpose of a schedular is to determine how much noise to add to the latent at a given step in the diffusion process. Let’s visualize the schedular function -

Code 
plt.plot(scheduler.sigmas)
plt.xlabel("Sampling step")
plt.ylabel("sigma")
plt.title("Schedular routine")
plt.show()

Fig. 9: Sampling schedule visualization.

The diffusion process follows this sampling schedule where we start with high noise and gradually denoise the image. Let’s visualize this process -

Code 
noise = torch.randn_like(latent_img) # Random noise
fig, axs = plt.subplots(2, 3, figsize=(16, 12))
for c, sampling_step in enumerate(range(0,51,10)):
    encoded_and_noised = scheduler.add_noise(latent_img, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
    axs[c//3][c%3].imshow(latents_to_pil(encoded_and_noised)[0])
    axs[c//3][c%3].set_title(f"Step - {sampling_step}")

Fig. 10: Noise progression through steps.

Let’s see how a U-Net removes the noise from the image. Let’s start by adding some noise to the image.

Code 
encoded_and_noised = scheduler.add_noise(latent_img, noise, timesteps=torch.tensor([scheduler.timesteps[40]]))
latents_to_pil(encoded_and_noised)[0]

Fig. 11: Noised Input fed to the U-Net.

Let’s run through U-Net and try to de-noise this image.

## Unconditional textual prompt
prompt = [""]

## Using clip model to get embeddings
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")

with torch.no_grad(): 
    text_embeddings = text_encoder(text_input.input_ids.to("cuda"))[0]
    
## Using U-Net to predict noise    
latent_model_input = torch.cat([encoded_and_noised.to("cuda").float()]).half()
with torch.no_grad():
    noise_pred = unet(latent_model_input, 40, encoder_hidden_states=text_embeddings)["sample"]

## Visualize after subtracting noise 
latents_to_pil(encoded_and_noised- noise_pred)[0]

Fig. 12: De-Noised Output from U-Net

As we can see above the U-Net output is clearer than the original noisy input passed.

4.3 What’s their role in the Stable diffusion pipeline

Latent diffusion uses the U-Net to gradually subtract noise in the latent space over several steps to reach the desired output. With each step, the amount of noise added to the latents is reduced till we reach the final de-noised output. U-Nets were first introduced by this paper for Biomedical image segmentation. The U-Net has an encoder and a decoder which are comprised of ResNet blocks. The stable diffusion U-Net also has cross-attention layers to provide them with the ability to condition the output based on the text description provided. The Cross-attention layers are added to both the encoder and the decoder part of the U-Net usually between ResNet blocks. You can learn more about this U-Net architecture here.

5 Conclusion

In this post, we saw the key components of a Stable diffusion pipeline i.e., CLIP Text encoder, VAE, and U-Net. In the next post, we will look at the diffusion process using these components. Read the next part here.

I hope you enjoyed reading it, and feel free to use my code and try it out for generating your images. Also, if there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn or email me at aayushmnit@gmail.com.

6 References

]]> Stable Diffusion https://aayushmnit.com/posts/2022-11-05-StableDiffusionP2/2022-11-05-StableDiffusionP2.html Sat, 05 Nov 2022 07:00:00 GMT Stable diffusion using 🤗 Hugging Face - Introduction Aayush Agrawal https://aayushmnit.com/posts/2022-11-02-StabeDiffusionP1/2022-11-02-StableDiffusionP1.html

A brief introduction to start generating images from text prompts using 🤗 hugging face - Diffusers library.

This is my first post of the Stable diffusion series, which I will write on Stable diffusion and other ongoing research happening in this field. Most of my learning can be attributed to knowledge acquired while doing the ‘From Deep learning foundations to Stable Diffusion’ course by FastAI and supplementing this with my research. The first few lessons of the FastAI course are publicly available here, and the rest will become available in early 2023. In this post, I want to give a brief introduction of how to use setup the 🤗 diffusion library and start generating images on your own. Next post, we will do a deep dive into mid-level components of this library.

1 Introduction

Stable diffusion simply put is a deep learning model which can generate an image given a textual prompt.

Fig. 1: Stable diffusion overview

As we can see from the image above we can pass a textual prompt like “A dog wearing a hat” and a stable diffusion model can generate an image representative of the text. Pretty amazing!

2 Using Hugging face Diffuser library

As with any python library, we need to follow certain installation steps before we can run it, here is a rundown of these steps.

2.1 Accepting the license

Before using the model, you need to go here and log in using your Hugging face account and then accept the model license to download and use the weights.

2.2 Token generation

If this is your first time using the hugging face library this might sound like a weird step. You need to go here and generate a token (preferably with write access) to download the model.

Fig. 2: Access token page

Once you have generated the token copy it. First, we will download the hugging face hub library using the following code.

!pip install huggingface-hub==0.10.1
Requirement already satisfied: huggingface-hub==0.10.1 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (0.10.1)
Requirement already satisfied: filelock in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from huggingface-hub==0.10.1) (3.8.0)
Requirement already satisfied: tqdm in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from huggingface-hub==0.10.1) (4.64.1)
Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from huggingface-hub==0.10.1) (4.3.0)
Requirement already satisfied: packaging>=20.9 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from huggingface-hub==0.10.1) (21.3)
Requirement already satisfied: requests in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from huggingface-hub==0.10.1) (2.28.1)
Requirement already satisfied: pyyaml>=5.1 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from huggingface-hub==0.10.1) (6.0)
Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from packaging>=20.9->huggingface-hub==0.10.1) (3.0.9)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from requests->huggingface-hub==0.10.1) (1.26.12)
Requirement already satisfied: certifi>=2017.4.17 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from requests->huggingface-hub==0.10.1) (2022.9.24)
Requirement already satisfied: charset-normalizer<3,>=2 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from requests->huggingface-hub==0.10.1) (2.0.4)
Requirement already satisfied: idna<4,>=2.5 in /home/aayush/miniconda3/envs/fastai/lib/python3.9/site-packages (from requests->huggingface-hub==0.10.1) (3.4)

Then use the following code, once you run it a widget will appear, paste your newly generated token and click login.

from huggingface_hub import notebook_login
notebook_login()
Login successful

Your token has been saved to /home/aayush/.huggingface/token

Authenticated through git-credential store but this isn't the helper defined on your machine.

You might have to re-authenticate when pushing to the Hugging Face Hub. Run the following command in your terminal in case you want to set this credential helper as the default



git config --global credential.helper store

2.3 Installing diffuser and transformer library

Once this process is done, install the dependencies using the following code. This will download the latest version of the diffusers and transformers library.

!pip install -qq -U diffusers transformers

That’s it, now we are ready to use the diffusers library.

3 Running Stable Diffusion

The first step is to import the StableDiffusionPipeline from the diffusers library.

from diffusers import StableDiffusionPipeline

The next step is to initialize a pipeline to generate an image. The first time you run the following command, it will download the model from the hugging face model hub to your local machine. You will require a GPU machine to be able to run this code.

pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4').to('cuda')

Now let’s pass a textual prompt and generate an image.

# Initialize a prompt
prompt = "a dog wearing hat"
# Pass the prompt in the pipeline
pipe(prompt).images[0]

Fig 3 - An example of image generated by the diffuser pipeline.

For further information on the diffusion pipeline read the documentation here.

4 Conclusion

In this post, we saw how to install diffusers library from hugging face and use the Stable diffusion model to generate images using a textual prompt. Read the part 2 here.

I hope you enjoyed reading it, and feel free to use my code and try it out for generating your images. Also, if there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn or email me at aayushmnit@gmail.com.

]]> Stable Diffusion https://aayushmnit.com/posts/2022-11-02-StabeDiffusionP1/2022-11-02-StableDiffusionP1.html Wed, 02 Nov 2022 07:00:00 GMT Model calibration for classification tasks using Python Aayush Agrawal https://aayushmnit.com/posts/2022-10-26-Model_Calibration/Model Calibration.html

A hands-on introduction to model calibration using Python.

What is Model Calibration?

When working with classification problems, machine learning models often produce a probabilistic outcome ranging between 0 to 1. This probabilistic output is then used by people to make decisions. Unfortunately, many machine learning models’ probabilistic outputs cannot be directly interpreted as the probability of an event happening. To achieve this outcome, the model needs to be calibrated.

Formally, a model is said to be perfectly calibrated if, for any probability value p, a prediction of a class with confidence p is correct 100*p percent of the time. In more simple terms, the probabilistic output equals the probability of occurrence.

Let us visualize a perfect calibrated and non-calibrated curve.

Code 
import numpy as np
import pandas as pd
import plotly.io as pio
import plotly.graph_objects as go
import plotly.express as px
pio.renderers.default = "plotly_mimetype+notebook_connected"
Code 
chart_df = pd.DataFrame({
    "actual_prob": np.arange(0,1.1, 0.1),
    "Calibrated": np.arange(0,1.1, 0.1),
    "Non-Calibrated":  [min(1,val - np.random.randn()*(1-idx)/20) for idx, val in enumerate(np.arange(0,1.1, 0.1))]
})
fig = px.line(
        data_frame = chart_df, 
        markers = True,
        x = "actual_prob", 
        y = ["Calibrated", "Non-Calibrated"],
        template = "plotly_dark")

fig.update_layout(
        title  = {"text": "Calibrated vs Non-calibrated model", "y": 0.98, "x": 0.5},
        xaxis_title="Predicted Probability",
        yaxis_title="Actual Probability",
        font = dict(size=15), 
        legend=dict(yanchor="bottom", y=0.95, orientation="h")
)

fig.update_traces(patch={"line": {"dash": "dash"}}) 
fig.show(renderer="notebook")

Fig 1 - A visualization of calibrated and non-calibrated curve.

On the x-axis, we have model output p which is between 0 and 1 and on the y-axis, we have fractions of positive captured within the predicted probability bin. We expect a linear relationship with slope 1.

Why do we need Model calibration?

There are many cases where model calibration is not required like in the case of ranking or selecting the top 20% for some targeting campaign. Calibrated models are especially important in making decision between multiple options of different magnitude or sizes, like expected value problems. In complex machine-learning decision engines, a machine-learning model might be used in conjunction with other machine-learning models.

For example, the business would like to prioritize upselling customer an additional product. So, a data scientist builds two ML models to estimate the probability of a customer buying two different products in addition to the products they are already purchasing.

Fig. 2: Upsell probability

In this case, Product 1 with $10 in revenue has an 80% probability of upselling, while Product 2 with $100 in revenue has a 60% chance of upselling. As a business, you might want to recommend Product 2 to the customer because of the higher expected value.

To do this expected value comparison between models, you need your models to be calibrated to output accurate probabilities.

Case study: Three methods for calibration

For demonstration in this article, we will be solving a binary classification problem using the Adult Income Dataset from the UCI machine learning Repository1 to predict if a certain individual income based on various census information (education level, age, gender, occupation, and more) exceeds $50K/year. For this example, we will limit our scope to the United States and the following predictors:

  • Age: continuous variable, individuals’ age
  • Occupation: categorical variable, Tech-support, Craft-repair, Other-service, Sales, Exec-managerial, Prof-specialty, Handlers-cleaners, Machine-op-inspct, Adm-clerical, Farming-fishing, Transport-moving, Priv-house-serv, Protective-serv, Armed-Forces.
  • HoursPerWeek: continuous variable, amount of hours spent in a job per week
  • Education: categorical variable, Bachelors, Some-college, 11th, HS-grad, Prof-school, Assoc-acdm, Assoc-voc, 9th, 7th-8th, 12th, Masters, 1st-4th, 10th, Doctorate, 5th-6th, Preschool.

In the below code we:

  • Import the libraries needed for this project
  • Load the data from our source
  • Label our columns
  • Filter data to only our desired columns, and limit our scope just to the United States for simplicity
  • Label encode our categorical columns
  • Set our random seed (for reproducibility)
  • Split our dataset into test and train
Code 
## Importing required libraries
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn import metrics
from sklearn.isotonic import IsotonicRegression
from sklearn.preprocessing import LabelEncoder
from sklearn.calibration import CalibratedClassifierCV
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
import matplotlib.pyplot as plt
%matplotlib inline
plt.style.use('dark_background')

def calibration_data(y_true, y_pred):
    df = pd.DataFrame({'y_true':y_true, 'y_pred_bucket': (y_pred//0.05)*0.05 + 0.025})
    cdf = df.groupby(['y_pred_bucket'], as_index=False).agg({'y_true':["mean","count"]})
    return cdf.y_true.values[:,0][cdf.y_true.values[:,1]>10], cdf.y_pred_bucket.values[cdf.y_true.values[:,1]>10]

def label_encoder(df,columns):
    '''
    Function to label encode
    Required Input - 
        - df = Pandas DataFrame
        - columns = List input of all the columns which needs to be label encoded
    Expected Output -
        - df = Pandas DataFrame with lable encoded columns
        - le_dict = Dictionary of all the column and their label encoders
    '''
    le_dict = {}
    for c in columns:
        print("Label encoding column - {0}".format(c))
        lbl = LabelEncoder()
        lbl.fit(list(df[c].values.astype('str')))
        df[c] = lbl.transform(list(df[c].values.astype('str')))
        le_dict[c] = lbl
    return df, le_dict
df = pd.read_csv( "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data", header=None)
df.columns = [
    "Age", "WorkClass", "fnlwgt", "Education", "EducationNum",
    "MaritalStatus", "Occupation", "Relationship", "Race", "Gender",
    "CapitalGain", "CapitalLoss", "HoursPerWeek", "NativeCountry", "Income"
]

## Filtering for Unites states
df = df.loc[df.NativeCountry == ' United-States',:]

## Only - Taking required columns
df = df.loc[:,["Education", "Age","Occupation", "HoursPerWeek", "Income"]]
df, _ = label_encoder(df, columns = ["Education", "Occupation", "Income"])
X = df.loc[:,["Education","Age", "Occupation", "HoursPerWeek"]]
y = df.Income

seed = 42
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20, random_state=seed)
Label encoding column - Education
Label encoding column - Occupation
Label encoding column - Income

Next, we:

  • Fit a random forest classifier on this data
  • Visualize the calibration plot
## Fitting the model on training data
rf_model = RandomForestClassifier(random_state=seed)
rf_model.fit(X_train, y_train)

## getting the output to visualize on test data
prob_true, prob_pred  = calibration_data(y_true = y_test, 
                                          y_pred = rf_model.predict_proba(X_test)[:,1])
Code 
chart_df = pd.DataFrame({
    "actuals": prob_true,
    "predicted": prob_pred,
    "expected": prob_pred
})
fig = px.line(
        data_frame = chart_df, 
        markers = True,
        x = "predicted", 
        y = ["actuals", "expected"], 
        template = "plotly_dark")
fig.update_layout(
        title  = {"text": "Calibration Plot: Without calibration", "y": 0.95, "x": 0.5},
        xaxis_title="Predicted Probability",
        yaxis_title="Actual Probability",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig 3 - An example of non-calibrated classifier

Note

As we can see above, the model is over-predicting in lower deciles and under-predicting in higher deciles.

Calibration method 1: Isotonic Regression

Isotonic regression is a variation of ordinary least squares regression. Isotonic regression has the added constraints that the predicted values must always be increasing or decreasing, and the predicted values have to lie as close to the observation as possible. Isotonic regression is often used in situations where the relationship between the input and output variables is known to be monotonic, but the exact form of the relationship is not known. In our case, we want a monotonic behavior where the actual probabilities must always be increasing with increasing predicted probability.

Mathematically, we can write the Isotonic regression as following - Where-

  • is the isotonic function
  • is the predicted probability from the original classifier f
  • is the calibrated probability

and we want to optimize the following function -

Isotonic regressions is prone to overfitting, hence it works better when there is enough training data available.

We will be using CalibratedClassifierCVto calibrate our classifier. The isotonic method fits a non-parametric isotonic regressor, which outputs a step-wise non-decreasing function. 2

## training model using random forest and isotonic regression for calibration
calibrated_rf = CalibratedClassifierCV(RandomForestClassifier(random_state=seed), method = 'isotonic')
calibrated_rf.fit(X_train, y_train)

## getting the output to visualize on test data
prob_true_calib, prob_pred_calib  = calibration_data(y_test, calibrated_rf.predict_proba(X_test)[:,1])
Code 
chart_df = pd.DataFrame({
    "actuals": prob_true_calib,
    "predicted": prob_pred_calib,
    "expected": prob_pred_calib
})
fig = px.line(
        data_frame = chart_df, 
        markers = True,
        x = "predicted", 
        y = ["actuals", "expected"], 
        template = "plotly_dark")
fig.update_layout(
        title  = {"text": "Calibration Plot: Isotonic", "y": 0.95, "x": 0.5},
        xaxis_title="Predicted Probability",
        yaxis_title="Actual Probability",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig 4 - An example of calibration using Isotonic regression

Note

As we can see above, the model is well calibrated once Isotonic regression is used because the actual probability lies much closer to the expected probability.

Calibration method 2: Sigmoid method / Platt scaling

The sigmoid method, also known as Platt scaling, works by transforming the outputs of the classifier using a sigmoid function. The calibrated probabilities are obtained using the following sigmoid function -

Where -

  • is the calibrated probability
  • is the predicted probability from the original classifier f
  • A and B are scalar parameters that learned by the algorithm

The parameters A and B are estimated using a maximum likelihood method that optimizes on the same training set as that for the original classifier f. 

Sigmoid/Platt scaling was originally invented for SVMs and it works well for other methods as well. This method is less prone to overfitting and should be preferred over Isotonic regression if you have less training data.

## training model using random forest and sigmoid method for calibration
calibrated_sigmoid = CalibratedClassifierCV(
     RandomForestClassifier(random_state=seed), method = 'sigmoid')
calibrated_sigmoid.fit(X_train, y_train)

## getting the output to visualize on test data
prob_true_calib, prob_pred_calib  = calibration_data(y_test, calibrated_sigmoid.predict_proba(X_test)[:,1])
Code 
chart_df = pd.DataFrame({
    "actuals": prob_true_calib,
    "predicted": prob_pred_calib,
    "expected": prob_pred_calib
})
fig = px.line(
        data_frame = chart_df, 
        markers = True,
        x = "predicted", 
        y = ["actuals", "expected"], 
        template = "plotly_dark")
fig.update_layout(
        title  = {"text": "Calibration Plot : Sigmoid", "y": 0.95, "x": 0.5},
        xaxis_title="Predicted Probability",
        yaxis_title="Actual Probability",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig 5 - An example of calibration using sigmoid method

Note

As we can see above, the model is more calibrated than our default model but not as well as the isotonic method.

Calibration method 3: Train with “Logloss” metric

Many boosting based ensembling classifier models like GradientBoostingClassifier, LightGBM, and Xgboost use log_loss as their default loss function. Training with log_loss helps the output probabilities to be calibrated. To explain why log loss helps in calibration, let’s look at binary log loss function -

Now let’s take an instance of observation where the true outcome is 1. Consider two predictions of 0.7 and 0.9 from the model. If we take the cutoff > 0.5 both of these predictions are correct but the logloss for 0.7 and 0.9 prediction is 0.36 and 0.1 respectively. As we can see log loss penalizes uncertainty in prediction forcing the model to predict close to the actual outcome.

Training with logloss is preferred over fitting a calibration function as model is already calibrated and reduces overhead of training and maintaining an extra model for calibration.

## Fitting the model on training data
gb_model = GradientBoostingClassifier()
gb_model.fit(X_train, y_train)

## getting the output to visualize on test data
prob_true, prob_pred  = calibration_data(y_true = y_test, 
                                          y_pred = gb_model.predict_proba(X_test)[:,1])
Code 
chart_df = pd.DataFrame({
    "actuals": prob_true,
    "predicted": prob_pred,
    "expected": prob_pred
})
fig = px.line(
        data_frame = chart_df, 
        markers = True,
        x = "predicted", 
        y = ["actuals", "expected"], 
        template = "plotly_dark")
fig.update_layout(
        title  = {"text": "Calibration Plot: Using boosting method", "y": 0.95, "x": 0.5},
        xaxis_title="Predicted Probability",
        yaxis_title="Actual Probability",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig 6 - An example of calibration using boosting method

Note

As we can see above, the model seems to be well calibrated. This is usually the case with most of the boosted methods as they use the log_loss function as their default loss function for binary classification.

Conclusion and Practical Guidance

In this article, we discussed what calibration is, in which applications it is important and why, and three different methods for calibration.

We demonstrated isotonic regression which can be used when there is enough training data available and model is not pre-calibrated.

We also demonstrated the Sigmoid method of calibration which works better when there is not enough training data and model is not pre-calibrated.

And finally, we demonstrated the Logloss metric for calibration which is the default in boosting methods and is our preferred method since the generated model is pre-calibrated and doesn’t require training and maintaining an extra model for calibration.

Footnotes

  1. Dua, D. and Graff, C. (2019). UCI Machine Learning Repository [http://archive.ics.uci.edu/ml]. Irvine, CA: University of California, School of Information and Computer Science. This dataset is licensed under a Creative Commons Attribution 4.0 International (CC BY 4.0) license.↩︎

  2. Scikit learn calibration documentation↩︎

]]> Model Calibration Machine Learning https://aayushmnit.com/posts/2022-10-26-Model_Calibration/Model Calibration.html Wed, 12 Oct 2022 07:00:00 GMT Mixing art into the science of model explainability Aayush Agrawal https://aayushmnit.com/posts/2022-09-23-Explainability/2022-09-23-Explainability.html

Overview on Explainable Boosting Machine and an approach for converting ML explanation to more human-friendly explanation.

Fig.1 - A lego figure on my desk

1. Science of ML explainability

1.1 The Interpretability vs Accuracy Trade-off

In traditional tabular machine learning approaches, Data scientists often deal with the trade-off b/w interpretability and accuracy.

Fig.2: Interpretability/Intelligibility and Accuracy Tradeoff
Image Credit - The Science Behind InterpretML: Explainable Boosting Machine

As shown in the chart above, we can see that Glass-Box models like Logistic Regression, Naive Bayes, and Decision Trees are simple models to interpret, and predictions from these models are not highly accurate. On the other hand, Black-Box models like Boosted Trees, Random Forest, and Neural Nets are hard to interpret but lead to highly accurate predictions.

1.2 Introducing EBMs

To solve the problem just mentioned above EBMs(Explainable Boosted Machine) model was developed by Microsoft Research1. “Explainable Boosting Machine (EBM) is a tree-based, cyclic gradient boosting Generalized Additive Model with automatic interaction detection. EBMs are often as accurate as state-of-the-art BlackBox models while remaining completely interpretable. Although EBMs are often slower to train than other modern algorithms, EBMs are extremely compact and fast at prediction time.”2

Fig.3: EBMs breaking the Interpretability vs Accuracy paradox
Image Credit - The Science Behind InterpretML: Explainable Boosting Machine

As we can see from the chart above, EBMs help us break out of this trade-off paradox and help us build models which are both highly interpretable and accurate. To further understand the math behind EBMs I highly encourage watching this 12-minute YouTube video -

Video - The Science Behind InterpretML: Explainable Boosting Machine

1.3 Glass box vs Black box models. What to choose?

Tip

The answer to every complex question in life is “It depends”.

There are trade-offs b/w using Glassbox models as compared to Blackbox models. There is no clear winner in picking one model over the other but depending on the situation DS can make an educated guess on what model to pick.

Fig.4: Glassbox models vs BlackBox models

Two considerations to think about while picking glass box vs black box models are the following-

1) Explainability Requirements - In the domain where there is no need for explanation or it is needed for a data scientist or technical audience for intuition/inspection purposes, in these cases, DS are well off using black box models. In the domain where an explanation is needed because of business or regulatory requirements or where these explanations are served to a non-technical audience (humans), glass-box models have an upper hand. This is because explanations coming out of the glass box models are exact and global.

Note

Exact and global just means that a value of a particular feature will always have the same effect on each prediction explanation. For example, in the case of the prediction of income of a particular individual being above $50k with age as one of the predictors, if the age is 40 and it will impact the target variable with the same proportion let us say 5% in each observation in the data where the age is 40. This is not the case when we build explanations through LIME and Shapely for black box models. In black-box models, age with the value 40 for example can have a 10% lift in an individual probability of their income being above 50k for one observation and -10% lift in the other.

2) Compute Requirement - DS needs to pay attention to various compute requirements for testing and training a model depending on its use case. EBMs are particularly slow in the training phase but provide fast predictions with built-in explanations. So, in cases where you need to train your model every hour, EBMs might not suffice your need. But, in cases where the training of the model happens monthly/weekly, and scores are generated on a more frequent basis (hourly/daily) EBMs might fit the use case well. Also, in cases where you might be required to produce an explanation for each prediction EBMs can save a lot of computing and might be the only feasible technique to use for millions of observations. Look below to understand the operational difference b/w EBMs and other tree-based ensemble methods.

Fig. 5: EBMs vs XgBoost/LightGBM

2. Hands on with EBMs

2.1 Data Overview

For this example, we will use Adult Income Dataset from the UCI machine learning Repository3. The problem in this dataset is set up as a binary classification problem to predict if a certain individual income based on various census information (education level, age, gender, occupation, etc.) exceeds $50K/year. For sake of simplicity, we are only going to use observations of individuals in the United States and the following predictors -

  • Age: continuous variable, individuals’ age
  • Occupation: categorical variable, Tech-support, Craft-repair, Other-service, Sales, Exec-managerial, Prof-specialty, Handlers-cleaners, Machine-op-inspct, Adm-clerical, Farming-fishing, Transport-moving, Priv-house-serv, Protective-serv, Armed-Forces.
  • HoursPerWeek: continuous variable, amount of hours spent in a job per week
  • Education: categorical variable, Bachelors, Some-college, 11th, HS-grad, Prof-school, Assoc-acdm, Assoc-voc, 9th, 7th-8th, 12th, Masters, 1st-4th, 10th, Doctorate, 5th-6th, Preschool.
Code 
## Importing required libraries
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn import metrics
from interpret.glassbox import ExplainableBoostingClassifier
from interpret import show
import warnings
import plotly.io as pio
import plotly.express as px
warnings.filterwarnings('ignore')
pio.renderers.default = "plotly_mimetype+notebook_connected"
## Loading the data
df = pd.read_csv( "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data", header=None)
df.columns = [
    "Age", "WorkClass", "fnlwgt", "Education", "EducationNum",
    "MaritalStatus", "Occupation", "Relationship", "Race", "Gender",
    "CapitalGain", "CapitalLoss", "HoursPerWeek", "NativeCountry", "Income"
]

## Filtering for Unites states
df = df.loc[df.NativeCountry == ' United-States',:]

## Only - Taking required columns
df = df.loc[:,["Education", "Age","Occupation", "HoursPerWeek", "Income"]]

df.head()
Education Age Occupation HoursPerWeek Income
0 Bachelors 39 Adm-clerical 40 <=50K
1 Bachelors 50 Exec-managerial 13 <=50K
2 HS-grad 38 Handlers-cleaners 40 <=50K
3 11th 53 Handlers-cleaners 40 <=50K
5 Masters 37 Exec-managerial 40 <=50K

Let’s look at target variable distribution.

Code 
plot_df = df.Income.value_counts().reset_index().rename(columns = {"index":"Income", "Income":"Count"})
fig = px.bar(plot_df, x = "Income", y = 'Count')
fig.update_layout(
        title  = {
            'text':"Target variable distribution",
            'y':0.95,
            'x':0.5,
        },
        legend =  dict(y=1, x= 0.8, orientation='v'),
        legend_title = "",
        xaxis_title="Income", 
        yaxis_title="Count of obersvations",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig 1 - Target variable distribution

Code 
print(df.Income.value_counts(normalize=True))
 <=50K    0.754165
 >50K     0.245835
Name: Income, dtype: float64

~24.6% of people in our dataset have income greater than $50K. The data looks good, we have the columns we need. We will use Education, Age, Occupation, and HoursPerWeek columns and predict Income. Before modeling, let us perform an 80-20 train-test split.

## Train-Test Split
X = df[df.columns[0:-1]]
y = df[df.columns[-1]]
seed = 1
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20, random_state=seed)
print(f"Data in training {len(y_train)}, Data in testing {len(y_test)}")
Data in training 23336, Data in testing 5834

2.2 Fitting an EBM Model

EBMs have a scikit-compatible API, so fitting the model and making predictions are the same as any scikit learn model.

ebm = ExplainableBoostingClassifier(random_state=seed, interactions=0)
ebm.fit(X_train, y_train)

auc = np.round(metrics.roc_auc_score((y_test != ' <=50K').astype(int).values, ebm.predict_proba(X_test)[:,1], ),3)
print(f"Accuracy: {np.round(np.mean(ebm.predict(X_test) == y_test)*100,2)}%, AUC: {auc}")
Accuracy: 80.12%, AUC: 0.828

I hope the above code block shows how similar the interpret-ml API is to the scikit learn API. Based on AUC on the validation set we can say our model is better than random predictions.

Tip

In practice, if you are dealing with millions of observations, Try doing feature selection using LightGBM/XGboost and only train your final models using EBMs. This will save you time in feature exploration.

2.3 Explaination from EBMs

Interpret package comes with both global and local explanations and has a variety of visualization tools to inspect what the model is learning.

2.3.1 Global explanations

Global explanations provide the following visualization -

  1. Summary - Feature importance plot, this chart provides the importance of each predictor in predicting the target variable
  2. Feature interaction with Prediction - This chart is the same look-up table EBM uses in making the actual prediction. This can help you in the inspection of how the feature value is contributing to prediction.
ebm_global = ebm.explain_global()
show(ebm_global, renderer='notebook')
Fig. 6: EBMs Global Explaination

2.3.2 Local explanations

The local explanation is our per-observation level explanation. EBMs have a great in-built visualization for displaying this information.

ebm_local = ebm.explain_local(X_test.iloc[0:5,:], y_test)
show(ebm_local, renderer='notebook')
Fig. 7: EBMs local Explaination

Let’s take one example of this explanation for observation at index 0 and look at it -

explainDF = pd.DataFrame.from_dict(
    {
        'names': ebm_local.data(0)['names'], 
        'data':ebm_local.data(0)['values'], 
        'contribution':ebm_local.data(0)['scores']
    })
explainDF
names data contribution
0 Education Bachelors 0.733420
1 Age 47 1.048227
2 Occupation ? -0.318846
3 HoursPerWeek 18 -0.854202

As we can see from the data, we can see we have the Name of the columns, the actual values, and the contribution of that value to the actual prediction score. For this observation let us see what the model is learning -

  1. Education as Bachelors is working in favor of >50K income
  2. Age value 47 is also in favor of >50K
  3. Occupation being “?” has a negative impact on >50K income
  4. HoursPerWeek being 18 has a negative impact on >50K income (Average work week hours in the US are around 40, so this makes sense)

You can also do it for the entire dataset and collect the importance of each feature. Here is a sample code to do the same.

scores = [x['scores'] for x in ebm_local._internal_obj['specific']]
summary = pd.DataFrame(scores)
summary.columns = ebm_local.data(0)['names']
summary.head()
Education Age Occupation HoursPerWeek
0 0.733420 1.048227 -0.318846 -0.854202
1 -0.990661 0.309251 0.171131 0.002109
2 -0.257254 0.735232 0.171131 0.002109
3 0.193118 0.682721 -0.417499 0.279677
4 0.733420 0.085672 0.389171 0.002109

Now we can extract the importance of all data rows in our test set.

2.4 Draw back and concerns

Credit - XKCD

These kinds of explanations are still very abstract, even at the observational level reasoning is not human(non-technical) friendly. When the feature count grows this becomes even non-human friendly. Typical business consumers of your model might not be well versed in reading such charts and shy away from trying the insights/predictions the model is giving them. After all, if I don’t understand something, I don’t trust it. That is where art comes in, let’s see how we can build on the above-derived observations and make it easier to understand.

3. The “Art” of ML explainability

Warning

The ideas I am going to share now are more marketing than real science.

To make people act on your model’s recommendation you must build trust. One idea is to build trust to support your explanations with data anecdotes.

Credit - Flickr

Data anecdotes exist everywhere, this idea came to my mind by looking at a stock market tool (image below). Notice two things they are highlighting -

  1. What is it? - The tool does an excellent job showing an event that has just happened. Example - “Microsoft Corp due to dividend announcement or “DELTA AIR LINES 14 Day RSI broke below the 70 level”.

  2. Why it matters? - Then the tool points to historical data and tells the significance of this event. In Microsoft’s example, when the event happens “Historically, the price of MSFT has risen an average of 11.9%”.

What if we can do the same with our models?

Fig 8. Snapshots taken from my fidelity tool

We can build a historical odds table from our training data. Let us try building one.

Code 
odds_data = X_train.copy()
odds_data['income'] = (y_train == " >50K").astype(int) 
## Converting continous variables in buckets
odds_data['AgeBucket'] =  (odds_data.Age // 5)
odds_data['HoursPerWeekBucket'] = (odds_data.HoursPerWeek // 5)

# Creating placeholder for odds dictionary
odds_dict = {} 

# Columns for which we need odds
columns = ['Education', 'AgeBucket', 'HoursPerWeekBucket', 'Occupation']
for colname in columns: #iterating through each column
    unique_val = odds_data[colname].unique() # Finding unique values in column
    ddict = {}
    for val in unique_val: # iterating each unique value in the column
        ## Odds that income is above > 50 in presence of the val
        val_p = odds_data.loc[odds_data[colname] == val, 'income'].mean() 
        ## Odds that income is above > 50 in absence of the val
        val_np = odds_data.loc[odds_data[colname] != val, 'income'].mean()
        
        ## Calculate lift
        if val_p >= val_np:
            odds = val_p / val_np
        else:
            odds = -1*val_np/(val_p+1e-3)
        
        ## Add to the col dict
        ddict[val] = np.round(odds,1)
    ## Add to the sub dict to odds dict
    odds_dict[colname] = ddict
print(odds_dict)
{'Education': {' HS-grad': -1.8, ' Some-college': -1.4, ' 10th': -4.0, ' Masters': 2.5, ' Assoc-voc': 1.1, ' Bachelors': 2.0, ' Assoc-acdm': 1.0, ' 1st-4th': -9.0, ' 11th': -4.7, ' 9th': -4.1, ' Doctorate': 3.1, ' Prof-school': 3.2, ' 12th': -2.9, ' 7th-8th': -3.8, ' 5th-6th': -6.2, ' Preschool': -245.4}, 'AgeBucket': {9: 1.8, 11: 1.5, 7: 1.4, 12: 1.2, 6: -1.1, 10: 1.8, 5: -2.3, 14: -1.2, 8: 1.6, 4: -20.0, 3: -100.0, 15: -1.3, 13: -1.1, 16: -2.1, 18: -2.4, 17: -245.3}, 'HoursPerWeekBucket': {8: -1.2, 9: 1.5, 12: 1.8, 3: -5.2, 7: -1.5, 6: -3.1, 5: -4.7, 10: 2.0, 4: -3.7, 14: 1.6, 2: -3.9, 13: 1.7, 11: 1.9, 15: 1.4, 1: -2.1, 0: -2.4, 19: 1.2, 16: 1.7, 17: 1.1, 18: 1.1}, 'Occupation': {' Machine-op-inspct': -1.8, ' ?': -2.4, ' Other-service': -7.1, ' Craft-repair': -1.0, ' Prof-specialty': 2.1, ' Handlers-cleaners': -3.9, ' Exec-managerial': 2.3, ' Sales': 1.1, ' Adm-clerical': -2.0, ' Transport-moving': -1.2, ' Tech-support': 1.2, ' Protective-serv': 1.4, ' Farming-fishing': -2.0, ' Priv-house-serv': -15.2, ' Armed-Forces': -1.9}}

Now using this odds table, we can generate predictions with the pre-filled template. Refer to the image below.

Fig 9. Output to human-readable text

With our explainDF data generated for row index 0 previously, we can use the above framework to convert it to text. Let us see what the output looks like -

Code 
def explainPredictions(df, pred, odds_dict):
    reasons = []
    if pred == 0:
        sdf =  df.loc[df.contribution < 0, :].sort_values(['contribution']).reset_index(drop=True).copy()
    else:
        sdf =  df.loc[df.contribution > 0, :].reset_index(drop=True).copy()
    
    for idx in range(sdf.shape[0]):
        col_name = sdf.names[idx]
        data = sdf.data[idx]
        if col_name in odds_dict:
            odd_value = odds_dict[col_name][data]
        else:
            odd_value = odds_dict[col_name+'Bucket'][data//5]
            
        s1 = f"This individual have {col_name} value '{data}'." 
        s2 = f"Historically, people with this behavior have {odd_value}x likely to have income over $50k."
        reasons.append(s1+s2)
    return reasons
explainPredictions(explainDF, ebm_local.data(0)['perf']['predicted'], odds_dict)
["This individual have HoursPerWeek value '18'.Historically, people with this behavior have -5.2x likely to have income over $50k.",
 "This individual have Occupation value ' ?'.Historically, people with this behavior have -2.4x likely to have income over $50k."]

What if this was a >50k prediction?

explainPredictions(explainDF, 1, odds_dict)
["This individual have Education value ' Bachelors'.Historically, people with this behavior have 2.0x likely to have income over $50k.",
 "This individual have Age value '47'.Historically, people with this behavior have 1.8x likely to have income over $50k."]

Looks great! We can generate human-readable recommendations for both >50K and <= 50K cases. Hope this provides you with some ideas on how to implement such human-readable explanations.

Summary

In the blog, we saw the interpretability-accuracy trade-off. How Explainable Boosting Machines (EBMs) work and how they help in building models which are highly interpretable and accurate. We saw how we can use EBMs through code to generate local and global explanations and how historical odd tables and some pre-generated text can help you convert these explanations into more human-friendly text.

I hope you enjoyed reading it, and feel free to use my approach to try it out for your purposes. Also, if there is any feedback on the code or just the blog post, feel free to reach out on LinkedIn or email me at aayushmnit@gmail.com. You can also follow me on Medium and Github for future blog posts and exploration project codes I might share.

Footnotes

  1. “InterpretML: A Unified Framework for Machine Learning Interpretability” (H. Nori, S. Jenkins, P. Koch, and R. Caruana 2019)↩︎

  2. “Interpret ML - EBM documentation”↩︎

  3. Dua, D. and Graff, C. (2019). UCI Machine Learning Repository [http://archive.ics.uci.edu/ml]. Irvine, CA: University of California, School of Information and Computer Science. This dataset is licensed under a Creative Commons Attribution 4.0 International (CC BY 4.0) license.↩︎

]]> Explainability Machine Learning https://aayushmnit.com/posts/2022-09-23-Explainability/2022-09-23-Explainability.html Fri, 23 Sep 2022 07:00:00 GMT Causal inference with Synthetic Control using Python and SparseSC Aayush Agrawal https://aayushmnit.com/posts/2022-09-19-SyntheticControl/2022-09-19_SyntheticControl.html

Understanding Synthetic Control and using Microsoft’s SparceSC package to run synthetic control on larger datasets.

What is Synthetic Control Method?

I will try to keep this part short and focus more on why Data scientists should care about such methods and how to use them on larger datasets based on practical experience using SparseSC package.

The Synthetic Control (SC) method is a statistical method used to estimate causal effects from binary treatments on observational panel (longitudinal) data. The method got quite a coverage by being described as “the most important innovation in the policy evaluation literature in the last few years” and got an article published in Washington Post - Seriously, here’s one amazing math trick to learn what can’t be known. “SC is a technique to create an artificial control group by taking a weighted average of untreated units in such a way that it reproduces the characteristics of the treated units before the intervention(treatment). The SC acts as the counterfactual for a treatment unit, and the estimate of a treatment effect is the difference between the observed outcome in the post-treatment period and the SC’s outcome.”1

“One way to think of SC is as an improvement upon difference-in-difference (DiD) estimation. Typical DiD will compare a treated unit to the average of the control units. But often the treated unit does not look like a typical control (e.g., it might have a different growth rate), in which case the ‘parallel trend’ assumption of DiD is not valid. SC remedies this by choosing a smarter linear combination, rather than the simple average, to weigh more heavily the more similar units. SC’s assumption is if there are endogenous factors that affect treatment and future outcomes then you should be able to control them by matching past outcomes. The matching that SC provides can therefore deal with some problems in estimation that DiD cannot handle.”2

Here is the link to the Causal inference book which I found most useful to understand the math behind SC- Causal Inference for The Brave and True by Matheus Facure - Chapter 15.

Why should any Data scientist care about this method?

Often as a Data Scientist, you will encounter situations as follows where running A/B testing is not feasible because of -

  1. Lack of infrastructure
  2. Lack of similar groups for running A/B testing (in case of evaluation of state policies, as there is no state equivalent of other)
  3. Providing unwanted advantage to one group over others. Sometimes running an A/B test can give an unfair advantage and lead you into anti-trust territory. For example, what if Amazon tries to charge differential pricing for different customers or apply different margins for their sellers for the same product?

As a data scientist, stakeholders may still ask you to estimate the impact of certain changes/treatments, and Synthetic controls can come to the rescue in this situation. For this reason, it is a valuable tool to keep in your algorithmic toolkit.

Problem Overview

We will use the Proposition 99 data to explain the use case for this approach and also how to use the SparceSC library and its key features. “In 1988, California passed a famous Tobacco Tax and Health Protection Act, which became known as Proposition 99. Its primary effect is to impose a 25-cent per pack state excise tax on the sale of tobacco cigarettes within California, with approximately equivalent excise taxes similarly imposed on the retail sale of other commercial tobacco products, such as cigars and chewing tobacco. Additional restrictions placed on the sale of tobacco include a ban on cigarette vending machines in public areas accessible by juveniles, and a ban on the individual sale of single cigarettes. Revenue generated by the act was earmarked for various environmental and health care programs, and anti-tobacco advertisements. To evaluate its effect, we can gather data on cigarette sales from multiple states and across a number of years. In our case, we got data from the year 1970 to 2000 from 39 states.”3

Code 
# Importing required libraries
import pandas as pd
import numpy as np
import SparseSC
from datetime import datetime
import warnings
import plotly.express as px
import plotly.graph_objects as pgo
pd.set_option("display.max_columns", None)
warnings.filterwarnings('ignore')

Let’s look at the data

#Import data
data_dir = "https://raw.githubusercontent.com/OscarEngelbrektson/SyntheticControlMethods/master/examples/datasets/"
df = pd.read_csv(data_dir + "smoking_data" + ".csv").drop(columns=["lnincome","beer", "age15to24"])
df.head()
state year cigsale retprice
0 Alabama 1970.0 89.8 39.6
1 Alabama 1971.0 95.4 42.7
2 Alabama 1972.0 101.1 42.3
3 Alabama 1973.0 102.9 42.1
4 Alabama 1974.0 108.2 43.1

We have data per state as treatment unit and yearly (year column) per-capita sales of cigarettes in packs (cigsale column) and the cigarette retail price (retprice column). We are going to pivot this data so that each row is one treatment unit(state), and columns represent the yearly cigsale value.

df = df.pivot(index= 'state', columns = 'year', values = "cigsale")
df.head()
year 1970.0 1971.0 1972.0 1973.0 1974.0 1975.0 1976.0 1977.0 1978.0 1979.0 1980.0 1981.0 1982.0 1983.0 1984.0 1985.0 1986.0 1987.0 1988.0 1989.0 1990.0 1991.0 1992.0 1993.0 1994.0 1995.0 1996.0 1997.0 1998.0 1999.0 2000.0
state
Alabama 89.8 95.4 101.1 102.9 108.2 111.7 116.2 117.1 123.0 121.4 123.2 119.6 119.1 116.3 113.0 114.5 116.3 114.0 112.1 105.6 108.6 107.9 109.1 108.5 107.1 102.6 101.4 104.9 106.2 100.7 96.2
Arkansas 100.3 104.1 103.9 108.0 109.7 114.8 119.1 122.6 127.3 126.5 131.8 128.7 127.4 128.0 123.1 125.8 126.0 122.3 121.5 118.3 113.1 116.8 126.0 113.8 108.8 113.0 110.7 108.7 109.5 104.8 99.4
California 123.0 121.0 123.5 124.4 126.7 127.1 128.0 126.4 126.1 121.9 120.2 118.6 115.4 110.8 104.8 102.8 99.7 97.5 90.1 82.4 77.8 68.7 67.5 63.4 58.6 56.4 54.5 53.8 52.3 47.2 41.6
Colorado 124.8 125.5 134.3 137.9 132.8 131.0 134.2 132.0 129.2 131.5 131.0 133.8 130.5 125.3 119.7 112.4 109.9 102.4 94.6 88.8 87.4 90.2 88.3 88.6 89.1 85.4 83.1 81.3 81.2 79.6 73.0
Connecticut 120.0 117.6 110.8 109.3 112.4 110.2 113.4 117.3 117.5 117.4 118.0 116.4 114.7 114.1 112.5 111.0 108.5 109.0 104.8 100.6 91.5 86.7 83.5 79.1 76.6 79.3 76.0 75.9 75.5 73.4 71.4

Let’s observe how cigarettes sales per capita is trending over time w.r.t California and other states.

Code 
plot_df = df.loc[df.index == "California"].T.reset_index(drop=False)
plot_df["OtherStates"] = df.loc[df.index != "California"].mean(axis=0).values


fig = px.line(
        data_frame = plot_df, 
        x = "year", 
        y = ["California","OtherStates"], 
        template = "plotly_dark")

fig.add_trace(
    pgo.Scatter(
        x=[1988,1988],
        y=[plot_df.California.min()*0.98,plot_df.OtherStates.max()*1.02], 
        line={
            'dash': 'dash',
        }, name='Proposition 99'
    ))
fig.update_layout(
        title  = {
            'text':"Gap in per-capita cigarette sales(in packs)",
            'y':0.95,
            'x':0.5,
        },
        legend =  dict(y=1, x= 0.8, orientation='v'),
        legend_title = "",
        xaxis_title="Year", 
        yaxis_title="Cigarette Sales Trend",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig 1 - Cigarette sales comparison b/w California and other states

As we can see from the chart above, we can see that there is a general decline in cigarette sales after the 1980s, and with the introduction of Proposition 99, the decreasing trend accelerated for the state of California. We cannot say for sure if this is happening with any statistical significance, it is just something we observed by examining the chart above.

To answer the question of whether Proposition 99 influenced cigarette consumption, we will use the pre-intervention period (1970-1988) to build a synthetic control group that mimics California cigarette sales trend. Then, we will see how this synthetic control behaves after the intervention.

Fitting Synthetic Control using SparseSC package

On a high level SparseSC package provide two functions for fitting Synthetic controls i.e., fit() method and fit_fast() method. On a high level -

  • fit() - This method tries to compute the weight jointly and results in SCs which are ‘optimal’. This is the most common method used in most of the code/libraries I have found but it is computationally expensive and takes a long time to run. Hence, does not scale for larger datasets.
  • fit_fast()- This method tries to compute the weight separately by performing some non-matching analysis. This solution is much faster and often the only feasible method with larger datasets. The authors of this package recommend the fit_fast method to start with and only move towards the fit method if needed.

The SparseSC.fit_fast() method required at least three arguments -

  1. features - This is the NumPy matrix of I/p variables where each row represents a treatment/control unit (states in our case), each column is the period from pre-treatment (1970-1988), and the value in the matrix is the metric of interest (in this case it is the cigsale value)
  2. targets - This is the NumPy matrix of I/p variables where each row represents a treatment/control unit (states in our case), each column is the period from post-treatment (1999-2000), and the value in the matrix is the metric of interest (in this case it is the cigsale value)
  3. treatment_units - This is the list of integers containing the row index value of treated units
Note

Note that treatment units can be a list of multiple treatment indexes. Think of cases where the same treatment is applied to multiple groups, for example, what if proposition 99 was rolled in both California and Minnesota State, in this case, treatment_units will get [2, 15], which are the respective index of these states.

## creating required features
features = df.iloc[:,df.columns <= 1988].values
targets = df.iloc[:,df.columns > 1988].values
treated_units = [idx for idx, val in enumerate(df.index.values) if val == 'California'] # [2]

## Fit fast model for fitting Synthetic controls
sc_model = SparseSC.fit_fast( 
    features=features,
    targets=targets,
    treated_units=treated_units
)

Now that we have fitted the model, let’s get the Synthetic Control output by using predict() function.

result = df.loc[df.index == 'California'].T.reset_index(drop=False)
result.columns = ["year", "Observed"] 
result['Synthetic'] = sc_model.predict(df.values)[treated_units,:][0]
result.head(5)
year Observed Synthetic
0 1970.0 123.0 122.394195
1 1971.0 121.0 125.114849
2 1972.0 123.5 129.704372
3 1973.0 124.4 126.753988
4 1974.0 126.7 126.276394

Now that we have our synthetic control, we can plot it with the outcome variable of the State of California.

Code 
fig = px.line(
        data_frame = result, 
        x = "year", 
        y = ["Observed","Synthetic"], 
        template = "plotly_dark",)

fig.add_trace(
    pgo.Scatter(
        x=[1988,1988],
        y=[result.Observed.min()*0.98,result.Observed.max()*1.02], 
        line={
            'dash': 'dash',
        }, name='Proposition 99'
    ))
fig.update_layout(
        title  = {
            'text':"Synthetic Control Assessment",
            'y':0.95,
            'x':0.5,
        },
        legend =  dict(y=1, x= 0.8, orientation='v'),
        legend_title = "",
        xaxis_title="Year", 
        yaxis_title="Per-capita cigarette sales (in packs)",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig - Assessment of Proposition 99 on state of California using Synthetic Control

Note

In the pre-intervention period, the synthetic control does not reproduce the treatment exactly but follows the curve closely. This is a good sign, as it indicates that we are not overfitting. Also, note that we do see divergence after the intervention (introduction of Proposition 99) after 1988.

With the synthetic control groups in hand, we can estimate the treatment effect as the gap between the treated and the synthetic control outcomes.

Code 
result['California Effect'] = result.Observed - result.Synthetic
fig = px.line(
        data_frame = result, 
        x = "year", 
        y = "California Effect", 
        template = "plotly_dark",)
fig.add_hline(0)
fig.add_trace(
    pgo.Scatter(
        x=[1988,1988],
        y=[result["California Effect"].min()*0.98,result["California Effect"].max()*1.02], 
        line={
            'dash': 'dash',
        }, name='Proposition 99'
    ))

fig.update_layout(
        title  = {
            'text':"Difference across time",
            'y':0.95,
            'x':0.5,
        },
        legend =  dict(y=1, x= 0.8, orientation='v'),
        legend_title = "",
        xaxis_title="Year", 
        yaxis_title="Gap in Per-capita cigarette sales (in packs)",
        font = dict(size=15)
)
fig.show(renderer='notebook')

Fig - Gap in Per-capita cigarette sales in California w.r.t Synthetic Control

Code 
print(f"Effect of Proposition 99 w.r.t Synthetic Control => {np.round(result.loc[result.year==2000,'California Effect'].values[0],1)} packs")
Effect of Proposition 99 w.r.t Synthetic Control => -28.8 packs

Looking at the chart above, we can observe that by the year 2000, Proposition 99 has reduced the sales of cigarettes by ~29 packs. Now we will figure out if this is statistically significant.

Making Inference

In Synthetic control, to find if the effect we observed is significant or not, we run a placebo test. A placebo test is taking a random untreated unit and pretending all other units are in control and fit the Synthetic control over this randomly selected untreated unit to estimate the effect. Once we have repeated this placebo test multiple times, we can estimate the distribution of this randomly observed effect and see if the effect observed is significantly different from the placebo observed effect. In the SparceSC package, we can use the estimate_effects method to do this automatically for us. The estimate effects method takes a minimum of two arguments -

  1. outcomes - This is the NumPy matrix of I/p variables where each row represents a treatment/control unit (states in our case), and each column is the period of both pre-treatment and post-treatment period and the value in the matrix is the metric of interest (in this case it is the cigsale value)
  2. unit_treatment_periods - Vector of treatment periods for each unit, (if a unit is never treated then use np.NaN if vector refers to periods by numerical index)
## Creating unit treatment_periods
unit_treatment_periods = np.full((df.values.shape[0]), np.nan)
unit_treatment_periods[treated_units] = [idx for idx, colname in enumerate(df.columns) if colname == 1988][0]

## fitting estimate effects method
sc = SparseSC.estimate_effects(
    outcomes = df.values,  
    unit_treatment_periods = unit_treatment_periods, 
    max_n_pl=50, # Number of placebos
    level=0.9 # Level for confidence intervals
)
print(sc)
Pre-period fit diagnostic: Were the treated harder to match in the pre-period than the controls were.
Average difference in outcome for pre-period between treated and SC unit (concerning if p-value close to 0 ): 
1.8073663793403947 (p-value: 0.9743589743589743)

(Investigate per-period match quality more using self.pl_res_pre.effect_vec)

Average Effect Estimation: -16.965374734951705 (p-value: 0.07692307692307693)

Effect Path Estimation:
 -2.712194133284143 (p-value: 0.6923076923076923)
-6.424223898557088 (p-value: 0.20512820512820512)
-4.18303325159971 (p-value: 0.5128205128205128)
-10.210104128966762 (p-value: 0.28205128205128205)
-10.198518971790051 (p-value: 0.2564102564102564)
-14.921163932323616 (p-value: 0.15384615384615385)
-18.441946863077938 (p-value: 0.1282051282051282)
-21.27793738802989 (p-value: 0.1794871794871795)
-22.47075245885469 (p-value: 0.1794871794871795)
-23.397467590595554 (p-value: 0.20512820512820512)
-26.13637789807555 (p-value: 0.05128205128205128)
-30.504443997992325 (p-value: 0.02564102564102564)
-29.67170704122487 (p-value: 0.05128205128205128)

The estimate_effects method returns an object which will print the treatment effect of each post-treatment year and estimate the significance of the observed difference. The information printed can also be found in the pl_res_pre function of the returned object.

print(f"Estimated effect of sales in California state in year 2000 because of preposition 99 is {np.round(sc.pl_res_post.effect_vec.effect[-1])}, \
with a p-value of  {np.round(sc.pl_res_post.effect_vec.p[-1],2)}")
Estimated effect of sales in California state in year 2000 because of preposition 99 is -30.0, with a p-value of  0.05

Conclusion

Here are some key takeaways -

  1. Synthetic control allows us to combine multiple control units to make them resemble the treated unit. With synthetic control, we can estimate what would have happened to our treated unit in the absence of treatment.
  2. Microsoft SparseSC library provides a fast and easy-to-use API to run synthetic control groups and allows us to run a placebo test to estimate the significance of the effects observed.
]]> Causal Inference Synthetic Control https://aayushmnit.com/posts/2022-09-19-SyntheticControl/2022-09-19_SyntheticControl.html Mon, 19 Sep 2022 07:00:00 GMT Finding similar images using Deep learning and Locality Sensitive Hashing Aayush Agrawal https://aayushmnit.com/posts/2019-03-17-Finding_similar_images_using_Deep_learning_and_Locality_Sensitive_Hashing/ Blog Transferred to Medium.com.

]]> Deep Learning FastAI Pytorch Vision https://aayushmnit.com/posts/2019-03-17-Finding_similar_images_using_Deep_learning_and_Locality_Sensitive_Hashing/ Sun, 17 Mar 2019 07:00:00 GMT Real-time Multi-Facial attribute detection using computer vision and deep learning with FastAI and OpenCV Aayush Agrawal https://aayushmnit.com/posts/2019-02-17-Multi_Facial_attribute_detection_using_FastAI_and_OpenCV/ Blog Transferred to Medium.com.

]]> Deep Learning FastAI Pytorch Vision Video https://aayushmnit.com/posts/2019-02-17-Multi_Facial_attribute_detection_using_FastAI_and_OpenCV/ Sun, 17 Feb 2019 08:00:00 GMT MultiLayer Perceptron using Fastai and Pytorch Aayush Agrawal https://aayushmnit.com/posts/2019-01-05-Multi_Layer_perceptron_using_Fastai_and_Pytorch/ Blog Transferred to Medium.com.

]]> Deep Learning FastAI Pytorch Vision https://aayushmnit.com/posts/2019-01-05-Multi_Layer_perceptron_using_Fastai_and_Pytorch/ Sat, 05 Jan 2019 08:00:00 GMT