{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc0e59cc",
   "metadata": {},
   "source": [
    "# Lesson 40: Document search with embeddings activity (solution)\n",
    "\n",
    "In this activity, you will build a simple document search system using word embeddings.\n",
    "\n",
    "1. **Document embeddings** - Create document vectors by averaging word embeddings\n",
    "2. **Similarity search** - Find documents similar to a query\n",
    "3. **Comparison** - Compare embedding-based search with TF-IDF search\n",
    "\n",
    "## Notebook set-up\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52374b61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gensim.downloader as api\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "# Load pre-trained GloVe embeddings (this may take a moment)\n",
    "print('Loading GloVe embeddings...')\n",
    "glove = api.load('glove-wiki-gigaword-100')\n",
    "print(f'Loaded {len(glove)} word vectors of dimension {glove.vector_size}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6980e0b1",
   "metadata": {},
   "source": [
    "## 1. Sample document corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d92d403",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample documents about different topics\n",
    "documents = [\n",
    "    'The quick brown fox jumps over the lazy dog',\n",
    "    'Machine learning algorithms can classify images and text',\n",
    "    'Neural networks are inspired by the human brain',\n",
    "    'Python is a popular programming language for data science',\n",
    "    'Deep learning has revolutionized computer vision',\n",
    "    'Natural language processing helps computers understand text',\n",
    "    'The cat sat on the mat near the window',\n",
    "    'Artificial intelligence is transforming many industries',\n",
    "    'Dogs and cats are popular household pets',\n",
    "    'Data scientists use statistics and machine learning',\n",
    "]\n",
    "\n",
    "print(f'Corpus contains {len(documents)} documents')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be0695c7",
   "metadata": {},
   "source": [
    "## 2. Create document embeddings\n",
    "\n",
    "### Task 1: Implement document embedding function\n",
    "\n",
    "Complete the function below to create a document embedding by averaging the word embeddings of all words in the document.\n",
    "\n",
    "**Hints:**\n",
    "- Tokenize the document using `.lower().split()`\n",
    "- Check if each word is in the embedding vocabulary using `word in glove`\n",
    "- Get the embedding with `glove[word]`\n",
    "- Average all valid word vectors using `np.mean(vectors, axis=0)`\n",
    "- Return a zero vector if no words are found"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "647e450c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_document_embedding(document: str, embeddings) -> np.ndarray:\n",
    "    '''Create document embedding by averaging word embeddings.'''\n",
    "    \n",
    "    # Step 1: Tokenize (lowercase and split)\n",
    "    words = document.lower().split()\n",
    "    \n",
    "    # Step 2: Get embeddings for words that exist in vocabulary\n",
    "    vectors = []\n",
    "\n",
    "    for word in words:\n",
    "\n",
    "        if word in embeddings:\n",
    "            vectors.append(embeddings[word])\n",
    "    \n",
    "    # Step 3: Return average or zero vector\n",
    "    if len(vectors) > 0:\n",
    "        return np.mean(vectors, axis=0)\n",
    "\n",
    "    else:\n",
    "        return np.zeros(embeddings.vector_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33d32e4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create embeddings for all documents\n",
    "doc_embeddings = np.array([get_document_embedding(doc, glove) for doc in documents])\n",
    "print(f'Document embeddings shape: {doc_embeddings.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24d33f7",
   "metadata": {},
   "source": [
    "## 3. Similarity search\n",
    "\n",
    "### Task 2: Implement document search function\n",
    "\n",
    "Complete the search function to find the most similar documents to a query.\n",
    "\n",
    "**Hints:**\n",
    "- Get the query embedding using `get_document_embedding`\n",
    "- Compute cosine similarity with `cosine_similarity`\n",
    "- Sort by similarity score and return top N results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b910b664",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_documents(query: str, documents: list, doc_embeddings: np.ndarray, top_n: int = 3):\n",
    "    '''Search for documents similar to query using embeddings.'''\n",
    "    \n",
    "    # Step 1: Get query embedding\n",
    "    query_embedding = get_document_embedding(query, glove)\n",
    "    \n",
    "    # Step 2: Compute cosine similarity with all documents\n",
    "    # cosine_similarity expects 2D arrays, so reshape query_embedding\n",
    "    similarities = cosine_similarity(\n",
    "        query_embedding.reshape(1, -1),\n",
    "        doc_embeddings\n",
    "    )\n",
    "    \n",
    "    # Step 3: Get indices sorted by similarity (descending)\n",
    "    sorted_indices = np.argsort(similarities[0])[::-1]\n",
    "    \n",
    "    # Step 4: Return top N results as list of (document, score) tuples\n",
    "    results = []\n",
    "\n",
    "    for idx in sorted_indices[:top_n]:\n",
    "        results.append((documents[idx], similarities[0][idx]))\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "761ad4f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the search function\n",
    "queries = [\n",
    "    'artificial intelligence and deep neural networks',\n",
    "    'pets and animals',\n",
    "    'programming code'\n",
    "]\n",
    "\n",
    "print('Embedding-based search results:\\n')\n",
    "\n",
    "for query in queries:\n",
    "\n",
    "    print(f'Query: \"{query}\"')\n",
    "    results = search_documents(query, documents, doc_embeddings)\n",
    "\n",
    "    for doc, score in results:\n",
    "        print(f'  [{score:.3f}] {doc}')\n",
    "\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4c84238",
   "metadata": {},
   "source": [
    "## 4. Comparison with TF-IDF search\n",
    "\n",
    "### Task 3: Compare embedding search with TF-IDF search\n",
    "\n",
    "Run the TF-IDF search below and compare results with embedding-based search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26f21954",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TF-IDF based search for comparison\n",
    "tfidf = TfidfVectorizer()\n",
    "tfidf_matrix = tfidf.fit_transform(documents)\n",
    "\n",
    "def search_tfidf(query: str, top_n: int = 3):\n",
    "    '''Search using TF-IDF similarity.'''\n",
    "\n",
    "    query_vec = tfidf.transform([query])\n",
    "    similarities = cosine_similarity(query_vec, tfidf_matrix)[0]\n",
    "    sorted_indices = np.argsort(similarities)[::-1]\n",
    "    \n",
    "    results = []\n",
    "\n",
    "    for idx in sorted_indices[:top_n]:\n",
    "        results.append((documents[idx], similarities[idx]))\n",
    "\n",
    "    return results\n",
    "\n",
    "print('TF-IDF search results:\\n')\n",
    "\n",
    "for query in queries:\n",
    "\n",
    "    print(f'Query: \"{query}\"')\n",
    "    results = search_tfidf(query)\n",
    "\n",
    "    for doc, score in results:\n",
    "        print(f'  [{score:.3f}] {doc}')\n",
    "\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1bc5c0a",
   "metadata": {},
   "source": [
    "## 5. Analysis questions\n",
    "\n",
    "After completing the tasks above, answer these questions:\n",
    "\n",
    "1. Which search method (embeddings vs TF-IDF) handles synonyms and related concepts better? Give an example.\n",
    "2. What queries where embeddings outperform TF-IDF, and vice versa?\n",
    "3. What are the trade-offs between these two approaches in terms of speed, memory, and accuracy?\n",
    "4. How might you combine both approaches for better search results?"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}