Exam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5405e217",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([5, 768]),\n",
       " torch.Size([12, 768]),\n",
       " torch.Size([12, 768]),\n",
       " torch.Size([12, 768]),\n",
       " torch.Size([6, 768]),\n",
       " torch.Size([16, 768])]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "model = AutoModel.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "documents = [\n",
    "    \"Bert is awesome\",\n",
    "    \"I'm looking for a job at Huggingface\",\n",
    "    \"I'm selling a used car in good condition\",\n",
    "    \"Huggingface team is awesome\",\n",
    "    \"My car is outside\",\n",
    "    \"Hugginface team is awesome\"\n",
    "    \"Faiss consumes a lot of memory\",\n",
    "]\n",
    "\n",
    "vectors = [\n",
    "  # tokenize the document, return it as PyTorch tensors (vectors),\n",
    "  # and pass it onto the model\n",
    "  model(**tokenizer(document, return_tensors='pt'))[0].detach().squeeze()\n",
    "  for document in documents\n",
    "]\n",
    "\n",
    "[v.size() for v in vectors]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "509dfb41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([768]),\n",
       " torch.Size([768]),\n",
       " torch.Size([768]),\n",
       " torch.Size([768]),\n",
       " torch.Size([768]),\n",
       " torch.Size([768])]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "averaged_vectors = [torch.mean(vector, dim=0) for vector in vectors]\n",
    "\n",
    "[v.size() for v in averaged_vectors]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a72d2a9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(document: str) -> torch.Tensor:\n",
    "  tokens = tokenizer(document, return_tensors='pt')\n",
    "  vector = model(**tokens)[0].detach().squeeze()\n",
    "  return torch.mean(vector, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "59581ad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "import numpy as np\n",
    "\n",
    "index = faiss.IndexIDMap(faiss.IndexFlatIP(768)) # the size of our vector space\n",
    "# index all the documents, we need them as numpy arrays first\n",
    "index.add_with_ids(\n",
    "    np.array([t.numpy() for t in averaged_vectors]),\n",
    "    # the IDs will be 0 to len(documents)\n",
    "    np.array(range(0, len(documents)))\n",
    ")\n",
    "\n",
    "def search(query: str, k=1):\n",
    "  encoded_query = encode(query).unsqueeze(dim=0).numpy()\n",
    "  top_k = index.search(encoded_query, k)\n",
    "  scores = top_k[0][0]\n",
    "  results = [documents[_id] for _id in top_k[1][0]]\n",
    "  return list(zip(results, scores))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1b105af3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(\"I'm looking for a job at Huggingface\", 62.06402),\n",
       " (\"I'm selling a used car in good condition\", 51.57861)]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "documents[1]\n",
    "# => \"I'm looking for a job at Huggingface\"\n",
    "\n",
    "search(documents[1], k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01b131f8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}