01-SentenceSimilarity.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sentence similarity\n",
    "\n",
    "In this notebook we'll run a basic sentence similarity model. This models will be later used for semantic search."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installing a sample model\n",
    "\n",
    "We'll use [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2). It's a small and lightweight model that will be good for this particular showcase."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "model_name = 'sentence-transformers/all-MiniLM-L6-v2'\n",
    "\n",
    "model = AutoModel.from_pretrained(model_name)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vectorizing sample data\n",
    "\n",
    "To see if the model is working correctly, we'll try to encode sample data and then try to use another sentence to find the closest passages of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qdrant_client import QdrantClient\n",
    "\n",
    "client = QdrantClient(host='localhost', port=6333)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a vector collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from qdrant_client.models import Distance, VectorParams\n",
    "\n",
    "client.recreate_collection(\n",
    "    collection_name=\"sample_collection\",\n",
    "    vectors_config=VectorParams(size=384, distance=Distance.COSINE),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add sample data into the collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "UpdateResult(operation_id=0, status=<UpdateStatus.COMPLETED: 'completed'>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from qdrant_client.models import PointStruct\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def mean_pooling(model_output, attention_mask):\n",
    "    token_embeddings = model_output[0] #First element of model_output contains all token embeddings\n",
    "    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n",
    "    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n",
    "\n",
    "sentences = [\n",
    "    \"Historia Polski jest bardzo ciekawa. Mieszko I był pierwszym władcą Polski. W roku 966 przyjął chrzest. W ten sposób Polska stała się krajem chrześcijańskim.\",\n",
    "    \"Maria Curie-Skłodowska była wybitną polską naukowczynią. Otrzymała dwie Nagrody Nobla. Była pierwszą kobietą, która otrzymała tę nagrodę. Chemia była jej pasją. Odkryła pierwiastki promieniotwórcze.\",\n",
    "    \"Trygonometria jest działem matematyki. Zajmuje się ona badaniem zależności między bokami i kątami trójkątów. Trygonometria jest bardzo ważna w nawigacji. Dzięki niej możemy określić położenie statku na morzu.\",\n",
    "]\n",
    "\n",
    "encoded_sentences = tokenizer(sentences, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    model_output = model(**encoded_sentences)\n",
    "\n",
    "sentence_embeddings = mean_pooling(model_output, encoded_sentences['attention_mask'])\n",
    "sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)\n",
    "\n",
    "client.upsert(\n",
    "    collection_name=\"sample_collection\",\n",
    "    points=[\n",
    "        PointStruct(\n",
    "            id=idx,\n",
    "            vector=vector,\n",
    "            payload={\"text\": sentences[idx]}\n",
    "        )\n",
    "        for idx, vector in enumerate(sentence_embeddings)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Search for similar vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best answer:\n",
      "Trygonometria jest działem matematyki. Zajmuje się ona badaniem zależności między bokami i kątami trójkątów. Trygonometria jest bardzo ważna w nawigacji. Dzięki niej możemy określić położenie statku na morzu.\n",
      "Similarity: 0.61781\n"
     ]
    }
   ],
   "source": [
    "query = \"Trójkąty są bardzo ważne w matematyce. Wszystkie figury geometryczne są bardzo ciekawe. Matematyka jest fascynująca.\"\n",
    "\n",
    "encoded_query = tokenizer(query, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    model_output = model(**encoded_query)\n",
    "\n",
    "query_embedding = mean_pooling(model_output, encoded_query['attention_mask'])\n",
    "query_embedding = F.normalize(query_embedding, p=2, dim=1)\n",
    "\n",
    "hits = client.search(\n",
    "    collection_name=\"sample_collection\",\n",
    "    query_vector=query_embedding[0].tolist(),\n",
    "    limit=1\n",
    ")\n",
    "\n",
    "best_answer = hits[0]\n",
    "\n",
    "print(\"Best answer:\")\n",
    "print(best_answer.payload['text'])\n",
    "print(\"Similarity:\", best_answer.score)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}