notebooks/en/code_search.ipynb (769 lines of code) (raw):
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "P9uMH4XjQQCh"
},
"source": [
"## Code Search with Vector Embeddings and Qdrant\n",
"\n",
"*Authored by: [Qdrant Team](https://qdrant.tech/)*"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "m2J2A8i2QQCi"
},
"source": [
"In this notebook, we demonstrate how you can use vector embeddings to navigate a codebase, and find relevant code snippets. We'll search codebases using natural semantic queries, and search for code based on a similar logic.\n",
"\n",
"You can check out the [live deployment](https://code-search.qdrant.tech/) of this approach which exposes the Qdrant codebase for search with a web interface."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GPXQ2PAzQQCi"
},
"source": [
"### The approach\n",
"\n",
"We need two models to accomplish our goal.\n",
"\n",
"- General usage neural encoder for Natural Language Processing (NLP), in our case [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2). We'll call this NLP model.\n",
"\n",
"- Specialized embeddings for code-to-code similarity search. We'll use the [jinaai/jina-embeddings-v2-base-code](https://huggingface.co/jinaai/jina-embeddings-v2-base-code) model for the task. It supports English and 30 widely used programming languages with a 8192 sequence length. Let's call this code model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZTKLHLfXQQCj"
},
"source": [
"To prepare our code for the NLP model, we need to preprocess the code to a format that closely resembles natural language. The code model supports a variety of standard programming languages, so there is no need to preprocess the snippets. We can use the code as is."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pLcwuG_bxZkg"
},
"source": [
"## Installing Dependencies\n",
"\n",
"Let's install the packages we'll work with.\n",
"\n",
"- [inflection](https://pypi.org/project/inflection/) - A string transformation library. It singularizes and pluralizes English words, and transforms CamelCase to underscored string.\n",
"- [fastembed](https://pypi.org/project/fastembed/) - A CPU-first, lightweight library for generating vector embeddings. [GPU support is available](https://github.com/qdrant/fastembed#%EF%B8%8F-fastembed-on-a-gpu).\n",
"- [qdrant-client](https://pypi.org/project/qdrant-client/) - Official Python library to interface with the Qdrant server."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "g31fBI58ykn9"
},
"outputs": [],
"source": [
"%pip install inflection qdrant-client fastembed"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xkYeQHOMQQCj"
},
"source": [
"### Data preparation\n",
"\n",
"Chunking the application sources into smaller parts is a non-trivial task. In general, functions, class methods, structs, enums, and all the other language-specific constructs are good candidates for chunks. They are big enough to contain some meaningful information, but small enough to be processed by embedding models with a limited context window. You can also use docstrings, comments, and other metadata can be used to enrich the chunks with additional information.\n",
"\n",
"<div style=\"text-align:center\"><img src=\"https://huggingface.co/datasets/Anush008/cookbook-images/resolve/main/data-chunking.png\" /></div>\n",
"\n",
"Text-based search is based on function signatures, but code search may return smaller pieces, such as loops. So, if we receive a particular function signature from the NLP model and part of its implementation from the code model, we merge the results."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EiwwXvDSQQCj"
},
"source": [
"### Parsing the Codebaase\n",
"\n",
"We'll use the [Qdrant codebase](https://github.com/qdrant/qdrant) for this demo.\n",
"While this codebase uses Rust, you can use this approach with any other language. You can use an [Language Server Protocol (LSP)](https://microsoft.github.io/language-server-protocol/) tool to build a graph of the codebase, and then extract chunks. We did our work with the [rust-analyzer](https://rust-analyzer.github.io/). We exported the parsed codebase into the [LSIF](https://microsoft.github.io/language-server-protocol/specifications/lsif/0.4.0/specification/) format, a standard for code intelligence data. Next, we used the LSIF data to navigate the codebase and extract the chunks."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can use the same approach for other languages. There are [plenty of implementations](https://microsoft.github.io/language-server-protocol/implementors/servers/) available."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i-F3D3xPQQCj"
},
"source": [
"We will then export the chunks into JSON documents with not only the code itself, but also context with the location of the code in the project."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SJwhq8qbQQCj"
},
"source": [
"You can examine the Qdrant structures, parsed in JSON, in the [structures.jsonl file](https://storage.googleapis.com/tutorial-attachments/code-search/structures.jsonl) in our Google Cloud Storage bucket. Download it and use it as a source of data for our code search."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Jc-kqwqLQQCj"
},
"outputs": [],
"source": [
"!wget https://storage.googleapis.com/tutorial-attachments/code-search/structures.jsonl"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "io57rPHNQQCk"
},
"source": [
"Next, load the file and parse the lines into a list of dictionaries:"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"id": "D0Z0vjTfQQCk"
},
"outputs": [],
"source": [
"import json\n",
"\n",
"structures = []\n",
"with open(\"structures.jsonl\", \"r\") as fp:\n",
" for i, row in enumerate(fp):\n",
" entry = json.loads(row)\n",
" structures.append(entry)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3A6pI3ElQQCk"
},
"source": [
"Let's see how one entry looks like."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "nZEBXstzQQCk",
"outputId": "81d723e6-08c2-4a25-8b8e-508c4a7e86b1"
},
"outputs": [],
"source": [
"structures[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```python\n",
"{'name': 'InvertedIndexRam',\n",
" 'signature': '# [doc = \" Inverted flatten index from dimension id to posting list\"] # [derive (Debug , Clone , PartialEq)] pub struct InvertedIndexRam { # [doc = \" Posting lists for each dimension flattened (dimension id -> posting list)\"] # [doc = \" Gaps are filled with empty posting lists\"] pub postings : Vec < PostingList > , # [doc = \" Number of unique indexed vectors\"] # [doc = \" pre-computed on build and upsert to avoid having to traverse the posting lists.\"] pub vector_count : usize , }',\n",
" 'code_type': 'Struct',\n",
" 'docstring': '= \" Inverted flatten index from dimension id to posting list\"',\n",
" 'line': 15,\n",
" 'line_from': 13,\n",
" 'line_to': 22,\n",
" 'context': {'module': 'inverted_index',\n",
" 'file_path': 'lib/sparse/src/index/inverted_index/inverted_index_ram.rs',\n",
" 'file_name': 'inverted_index_ram.rs',\n",
" 'struct_name': None,\n",
" 'snippet': '/// Inverted flatten index from dimension id to posting list\\n#[derive(Debug, Clone, PartialEq)]\\npub struct InvertedIndexRam {\\n /// Posting lists for each dimension flattened (dimension id -> posting list)\\n /// Gaps are filled with empty posting lists\\n pub postings: Vec<PostingList>,\\n /// Number of unique indexed vectors\\n /// pre-computed on build and upsert to avoid having to traverse the posting lists.\\n pub vector_count: usize,\\n}\\n'}}\n",
" ```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aaKCggNgQQCk"
},
"source": [
"### Code to natural language conversion\n",
"\n",
"Each programming language has its own syntax which is not a part of the natural language. Thus, a general-purpose model probably does not understand the code as is. We can, however, normalize the data by removing code specifics and including additional context, such as module, class, function, and file name. We take the following steps:\n",
"\n",
"1. Extract the signature of the function, method, or other code construct.\n",
"2. Divide camel case and snake case names into separate words.\n",
"3. Take the docstring, comments, and other important metadata.\n",
"4. Build a sentence from the extracted data using a predefined template.\n",
"5. Remove the special characters and replace them with spaces."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hbmUDCb9QQCk"
},
"source": [
"We can now define the `textify` function that uses the `inflection` library to carry out our conversions:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"id": "ubGOBozmQQCk"
},
"outputs": [],
"source": [
"import inflection\n",
"import re\n",
"\n",
"from typing import Dict, Any\n",
"\n",
"\n",
"def textify(chunk: Dict[str, Any]) -> str:\n",
" # Get rid of all the camel case / snake case\n",
" # - inflection.underscore changes the camel case to snake case\n",
" # - inflection.humanize converts the snake case to human readable form\n",
" name = inflection.humanize(inflection.underscore(chunk[\"name\"]))\n",
" signature = inflection.humanize(inflection.underscore(chunk[\"signature\"]))\n",
"\n",
" # Check if docstring is provided\n",
" docstring = \"\"\n",
" if chunk[\"docstring\"]:\n",
" docstring = f\"that does {chunk['docstring']} \"\n",
"\n",
" # Extract the location of that snippet of code\n",
" context = (\n",
" f\"module {chunk['context']['module']} \" f\"file {chunk['context']['file_name']}\"\n",
" )\n",
" if chunk[\"context\"][\"struct_name\"]:\n",
" struct_name = inflection.humanize(\n",
" inflection.underscore(chunk[\"context\"][\"struct_name\"])\n",
" )\n",
" context = f\"defined in struct {struct_name} {context}\"\n",
"\n",
" # Combine all the bits and pieces together\n",
" text_representation = (\n",
" f\"{chunk['code_type']} {name} \"\n",
" f\"{docstring}\"\n",
" f\"defined as {signature} \"\n",
" f\"{context}\"\n",
" )\n",
"\n",
" # Remove any special characters and concatenate the tokens\n",
" tokens = re.split(r\"\\W\", text_representation)\n",
" tokens = filter(lambda x: x, tokens)\n",
" return \" \".join(tokens)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3hlqx9MGQQCk"
},
"source": [
"Now we can use `textify` to convert all chunks into text representations:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"id": "E6VEQnuIQQCk"
},
"outputs": [],
"source": [
"text_representations = list(map(textify, structures))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "J_GlAXvXQQCl"
},
"source": [
"Let's see how one of our representations looks like:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 53
},
"id": "zosN7TC9QQCl",
"outputId": "38e0d938-3fb1-4426-f00c-74a42267bf7d"
},
"outputs": [],
"source": [
"text_representations[1000]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```python\n",
"'Function Hnsw discover precision that does Checks discovery search precision when using hnsw index this is different from the tests in defined as Fn hnsw discover precision module integration file hnsw_discover_test rs'\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZM4h4wQcQQCl"
},
"source": [
"### Natural language embeddings\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7gp7y_jYQQCl"
},
"outputs": [],
"source": [
"from fastembed import TextEmbedding\n",
"\n",
"batch_size = 5\n",
"\n",
"nlp_model = TextEmbedding(\"sentence-transformers/all-MiniLM-L6-v2\", threads=0)\n",
"nlp_embeddings = nlp_model.embed(text_representations, batch_size=batch_size)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LQZft9BqQQCl"
},
"source": [
"### Code Embeddings"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"id": "fPqV4sDJQQCl"
},
"outputs": [],
"source": [
"code_snippets = [structure[\"context\"][\"snippet\"] for structure in structures]\n",
"\n",
"code_model = TextEmbedding(\"jinaai/jina-embeddings-v2-base-code\")\n",
"\n",
"code_embeddings = code_model.embed(code_snippets, batch_size=batch_size)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sm4KSxVeQQCl"
},
"source": [
"### Building Qdrant collection"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JcU_BileQQCl"
},
"source": [
"Qdrant supports multiple modes of deployment. Including in-memory for prototyping, Docker and Qdrant Cloud. You can refer to the [installation instructions](https://qdrant.tech/documentation/guides/installation/) for more information.\n",
"\n",
"We'll continue the tutorial using an in-memory instance.\n",
"\n",
"> [!TIP]\n",
"> In-memory can only be used for quick-prototyping and tests. It is a Python implementation of the Qdrant server methods.\n",
"\n",
"Let's create a collection to store our vectors."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"id": "ZOLKqa7sQQCl"
},
"outputs": [],
"source": [
"from qdrant_client import QdrantClient, models\n",
"\n",
"COLLECTION_NAME = \"qdrant-sources\"\n",
"\n",
"client = QdrantClient(\":memory:\") # Use in-memory storage\n",
"# client = QdrantClient(\"http://locahost:6333\") # For Qdrant server\n",
"\n",
"client.create_collection(\n",
" COLLECTION_NAME,\n",
" vectors_config={\n",
" \"text\": models.VectorParams(\n",
" size=384,\n",
" distance=models.Distance.COSINE,\n",
" ),\n",
" \"code\": models.VectorParams(\n",
" size=768,\n",
" distance=models.Distance.COSINE,\n",
" ),\n",
" },\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ixffRmB0QQCl"
},
"source": [
"Our newly created collection is ready to accept the data. Let’s upload the embeddings:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ZOLKqa7sQQCl"
},
"outputs": [],
"source": [
"from tqdm import tqdm\n",
"\n",
"points = []\n",
"total = len(structures)\n",
"print(\"Number of points to upload: \", total)\n",
"\n",
"for id, (text_embedding, code_embedding, structure) in tqdm(enumerate(zip(nlp_embeddings, code_embeddings, structures)), total=total):\n",
" # FastEmbed returns generators. Embeddings are computed as consumed.\n",
" points.append(\n",
" models.PointStruct(\n",
" id=id,\n",
" vector={\n",
" \"text\": text_embedding,\n",
" \"code\": code_embedding,\n",
" },\n",
" payload=structure,\n",
" )\n",
" )\n",
"\n",
" # Upload points in batches\n",
" if len(points) >= batch_size:\n",
" client.upload_points(COLLECTION_NAME, points=points, wait=True)\n",
" points = []\n",
"\n",
"# Ensure any remaining points are uploaded\n",
"if points:\n",
" client.upload_points(COLLECTION_NAME, points=points)\n",
"\n",
"print(f\"Total points in collection: {client.count(COLLECTION_NAME).count}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UCRXvkJZQQCm"
},
"source": [
"The uploaded points are immediately available for search. Next, query the collection to find relevant code snippets."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3jBAVfB7QQCm"
},
"source": [
"### Querying the codebase\n",
"\n",
"We use one of the models to search the collection via Qdrant's new [Query API](https://qdrant.tech/blog/qdrant-1.10.x/). Start with text embeddings. Run the following query “How do I count points in a collection?”. Review the results."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"id": "cUoLe_SyQQCm"
},
"outputs": [],
"source": [
"query = \"How do I count points in a collection?\"\n",
"\n",
"hits = client.query_points(\n",
" COLLECTION_NAME,\n",
" query=next(nlp_model.query_embed(query)).tolist(),\n",
" using=\"text\",\n",
" limit=3,\n",
").points"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "k2vbMhW2QQCm"
},
"source": [
"Now, review the results. The following table lists the module, the file name\n",
"and score. Each line includes a link to the signature.\n",
"\n",
"| module | file_name | score | signature |\n",
"|--------------------|---------------------|------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n",
"| operations | types.rs | 0.5493385 | [`pub struct CountRequestInternal`](https://github.com/qdrant/qdrant/blob/4aac02315bb3ca461a29484094cf6d19025fce99/lib/collection/src/operations/types.rs#L794) |\n",
"| map_index | types.rs | 0.49973965 | [`fn get_points_with_value_count`](https://github.com/qdrant/qdrant/blob/4aac02315bb3ca461a29484094cf6d19025fce99/lib/segment/src/index/field_index/map_index/mod.rs#L89) |\n",
"| map_index | mutable_map_index.rs | 0.49941066 | [`pub fn get_points_with_value_count`](https://github.com/qdrant/qdrant/blob/4aac02315bb3ca461a29484094cf6d19025fce99/lib/segment/src/index/field_index/map_index/mutable_map_index.rs#L143) |\n",
"\n",
"It seems we were able to find some relevant code structures. Let's try the same with the code embeddings:"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {
"id": "Z2vczTXCQQCm"
},
"outputs": [],
"source": [
"hits = client.query_points(\n",
" COLLECTION_NAME,\n",
" query=next(code_model.query_embed(query)).tolist(),\n",
" using=\"code\",\n",
" limit=3,\n",
").points"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "24gl0XfIQQCp"
},
"source": [
"Output:\n",
"\n",
"| module | file_name | score | signature |\n",
"|---------------|----------------------------|------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n",
"| field_index | geo_index.rs | 0.7217579 | [`fn count_indexed_points`](https://github.com/qdrant/qdrant/blob/4aac02315bb3ca461a29484094cf6d19025fce99/lib/segment/src/index/field_index/geo_index/mod.rs#L319) |\n",
"| numeric_index | mod.rs | 0.7113214 | [`fn count_indexed_points`](https://github.com/qdrant/qdrant/blob/4aac02315bb3ca461a29484094cf6d19025fce99/lib/segment/src/index/field_index/numeric_index/mod.rs#L317) |\n",
"| full_text_index | text_index.rs | 0.6993165 | [`fn count_indexed_points`](https://github.com/qdrant/qdrant/blob/4aac02315bb3ca461a29484094cf6d19025fce99/lib/segment/src/index/field_index/full_text_index/text_index.rs#L179) |"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Uo4HFN1SQQCp"
},
"source": [
"While the scores retrieved by different models are not comparable, but we can\n",
"see that the results are different. Code and text embeddings can capture\n",
"different aspects of the codebase. We can use both models to query the collection\n",
"and then combine the results to get the most relevant code snippets."
]
},
{
"cell_type": "code",
"execution_count": 25,
"metadata": {
"id": "266Erpi2VdLP"
},
"outputs": [],
"source": [
"from qdrant_client import models\n",
"\n",
"hits = client.query_points(\n",
" collection_name=COLLECTION_NAME,\n",
" prefetch=[\n",
" models.Prefetch(\n",
" query=next(nlp_model.query_embed(query)).tolist(),\n",
" using=\"text\",\n",
" limit=5,\n",
" ),\n",
" models.Prefetch(\n",
" query=next(code_model.query_embed(query)).tolist(),\n",
" using=\"code\",\n",
" limit=5,\n",
" ),\n",
" ],\n",
" query=models.FusionQuery(fusion=models.Fusion.RRF)\n",
").points"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "6o9ubjlM46Y1",
"outputId": "5f58833a-4e6a-4b47-de94-f540fd605c9b"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"| operations | lib/collection/src/operations/types.rs | 0.5 | ` # [doc = \" Count Request\"] # [doc = \" Counts the number of points which satisfy the given filter.\"] # [doc = \" If filter is not provided, the count of all points in the collection will be returned.\"] # [derive (Debug , Deserialize , Serialize , JsonSchema , Validate)] # [serde (rename_all = \"snake_case\")] pub struct CountRequestInternal { # [doc = \" Look only for points which satisfies this conditions\"] # [validate] pub filter : Option < Filter > , # [doc = \" If true, count exact number of points. If false, count approximate number of points faster.\"] # [doc = \" Approximate count might be unreliable during the indexing process. Default: true\"] # [serde (default = \"default_exact_count\")] pub exact : bool , } ` |\n",
"| field_index | lib/segment/src/index/field_index/geo_index.rs | 0.5 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| map_index | lib/segment/src/index/field_index/map_index/mod.rs | 0.33333334 | ` fn get_points_with_value_count < Q > (& self , value : & Q) -> Option < usize > where Q : ? Sized , N : std :: borrow :: Borrow < Q > , Q : Hash + Eq , ` |\n",
"| numeric_index | lib/segment/src/index/field_index/numeric_index/mod.rs | 0.33333334 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| fixtures | lib/segment/src/fixtures/payload_context_fixture.rs | 0.25 | ` fn total_point_count (& self) -> usize ` |\n",
"| map_index | lib/segment/src/index/field_index/map_index/mutable_map_index.rs | 0.25 | ` fn get_points_with_value_count < Q > (& self , value : & Q) -> Option < usize > where Q : ? Sized , N : std :: borrow :: Borrow < Q > , Q : Hash + Eq , ` |\n",
"| id_tracker | lib/segment/src/id_tracker/simple_id_tracker.rs | 0.2 | ` fn total_point_count (& self) -> usize ` |\n",
"| map_index | lib/segment/src/index/field_index/map_index/mod.rs | 0.2 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| map_index | lib/segment/src/index/field_index/map_index/mod.rs | 0.16666667 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| field_index | lib/segment/src/index/field_index/stat_tools.rs | 0.16666667 | ` fn number_of_selected_points (points : usize , values : usize) -> usize ` |\n"
]
}
],
"source": [
"for hit in hits:\n",
" print(\n",
" \"| \",\n",
" hit.payload[\"context\"][\"module\"], \" | \",\n",
" hit.payload[\"context\"][\"file_path\"], \" | \",\n",
" hit.score, \" | `\",\n",
" hit.payload[\"signature\"], \"` |\"\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kjG6JKy7QQCp"
},
"source": [
"This is one example of how you can fuse the results from different models.\n",
"In a real-world scenario, you might run some reranking and deduplication, as well as additional processing of the results."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xFiW5eO3QQCp"
},
"source": [
"### Grouping the results\n",
"\n",
"You can improve the search results, by grouping them by payload properties.\n",
"In our case, we can group the results by the module. If we use code embeddings,\n",
"we can see multiple results from the `map_index` module. Let's group the\n",
"results and assume a single result per module:"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {
"id": "PVBkf-zrQQCp"
},
"outputs": [],
"source": [
"results = client.query_points_groups(\n",
" COLLECTION_NAME,\n",
" query=next(code_model.query_embed(query)).tolist(),\n",
" using=\"code\",\n",
" group_by=\"context.module\",\n",
" limit=5,\n",
" group_size=1,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "kwJRvi6X50vs",
"outputId": "20192503-6b5b-44c4-9aff-52a309cb75a8"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"| field_index | geo_index.rs | 0.7217579 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| numeric_index | mod.rs | 0.7113214 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| fixtures | payload_context_fixture.rs | 0.6993165 | ` fn total_point_count (& self) -> usize ` |\n",
"| map_index | mod.rs | 0.68385994 | ` fn count_indexed_points (& self) -> usize ` |\n",
"| full_text_index | text_index.rs | 0.6660142 | ` fn count_indexed_points (& self) -> usize ` |\n"
]
}
],
"source": [
"for group in results.groups:\n",
" for hit in group.hits:\n",
" print(\n",
" \"| \",\n",
" hit.payload[\"context\"][\"module\"], \" | \",\n",
" hit.payload[\"context\"][\"file_name\"], \" | \",\n",
" hit.score, \" | `\",\n",
" hit.payload[\"signature\"], \"` |\"\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mxiuYai_QQCq"
},
"source": [
"That concludes our tutorial. Thanks for taking the time to get here. We've just begun exploring what's possible with vector embeddings and how to improve it. Feel free to experiment your way; you could build something very cool! Do share it with us 🙏 We are [here](https://qdrant.tech/contact-us/)."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"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.11.8"
}
},
"nbformat": 4,
"nbformat_minor": 0
}