Basic RAG

Basic RAG

Retrieval-Augmented Generation (RAG) is a technique that combines the strengths of pre-trained language models with the ability to retrieve information from a large corpus of documents. RAG enables the language model to produce more informed, accurate, and contextually relevant answers than by relying on its pre-trained knowledge alone.

At Cohere, all RAG calls come with... precise citations! πŸŽ‰
The model cites which groups of words, in the RAG chunks, were used to generate the final answer.
These citations make it easy to check where the model’s generated response claims are coming from and they help users gain visibility into the model reasoning.

RAG consists of 3 steps:

  • Step 1: Indexing and given a user query, retrieve the relevant chunks from the index
  • Step 2: Optionally, rerank the retrieved chunks
  • Step 3: Generate the model final answer with precise citations, given the retrieved and reranked chunks

Step 0 - Imports & Getting some data

In this example, we'll use a recent piece of text, that wasn't in the training data: the Wikipedia page of the movie "Dune 2".

In practice, you would typically do RAG on much longer text, that doesn't fit in the context window of the model.

%pip install "cohere<5" --quiet
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 52.8/52.8 kB 1.6 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.1/3.1 MB 33.6 MB/s eta 0:00:00
[?25h
import cohere
API_KEY = "..." # fill in your Cohere API key here
co = cohere.Client(API_KEY)
!pip install wikipedia --quiet
import wikipedia
  Preparing metadata (setup.py) ... [?25l[?25hdone
  Building wheel for wikipedia (setup.py) ... [?25l[?25hdone
article = wikipedia.page('Dune Part Two')
text = article.content
print(f"The text has roughly {len(text.split())} words.")
The text has roughly 5323 words.

Step 1 - Indexing and given a user query, retrieve the relevant chunks from the index

We index the document in a vector database. This requires getting the documents, chunking them, embedding, and indexing them in a vector database. Then we retrieved relevant results based on the users' query.

We split the document into chunks of roughly 512 words

%pip install -qU langchain-text-splitters --quiet
from langchain_text_splitters import RecursiveCharacterTextSplitter
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 256.9/256.9 kB 6.6 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 66.6/66.6 kB 8.2 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 138.5/138.5 kB 14.5 MB/s eta 0:00:00
[?25h
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=50,
    length_function=len,
    is_separator_regex=False,
)

chunks_ = text_splitter.create_documents([text])
chunks = [c.page_content for c in chunks_]
print(f"The text has been broken down in {len(chunks)} chunks.")
The text has been broken down in 91 chunks.

Embed every text chunk

Cohere embeddings are state-of-the-art.

model="embed-english-v3.0"
response = co.embed(
    texts= chunks,
    model=model,
    input_type="search_document",
    embedding_types=['float']
)
embeddings = response.embeddings.float
print(f"We just computed {len(embeddings)} embeddings.")
We just computed 91 embeddings.

Store the embeddings in a vector database

We use the simplest vector database ever: a python dictionary using np.array().

!pip install numpy --quiet
import numpy as np
vector_database = {i: np.array(embedding) for i, embedding in enumerate(embeddings)}

Given a user query, retrieve the relevant chunks from the vector database

Define the user question

query = "Name everyone involved in writing the script, directing, and producing 'Dune: Part Two'?"

Embed the user question

Cohere embeddings are state-of-the-art.

response = co.embed(
    texts=[query],
    model=model,
    input_type="search_query",
    embedding_types=['float']
)
query_embedding = response.embeddings.float[0]
print("query_embedding: ", query_embedding)
query_embedding:  [-0.068603516, -0.02947998, -0.06274414, -0.015449524, -0.033294678, 0.0056877136, -0.047210693, 0.04714966, -0.024871826, 0.008148193, 0.0770874, 0.023880005, -0.058685303, -0.052520752, 0.012832642, 0.024398804, 0.0053215027, 0.035491943, 0.02961731, -0.0069847107, 0.01083374, -0.0011358261, -0.002199173, 0.018417358, 0.027389526, -0.002691269, -0.026535034, 0.015197754, 0.024368286, 0.03729248, 0.0057754517, -0.02229309, -0.014694214, 0.019989014, -0.0036315918, -0.013793945, 0.02835083, 0.006011963, 0.011428833, 0.008682251, 0.046142578, -0.040039062, -0.032196045, -0.002653122, -0.012580872, -0.0041618347, 0.03111267, -0.016799927, 0.014801025, -0.00030636787, -0.033050537, 0.033966064, -0.016021729, -0.025009155, -0.007534027, -0.017074585, 0.008415222, -0.10620117, 0.019195557, -0.015686035, -0.0043182373, -0.045440674, 0.05404663, 0.030776978, -0.014129639, -0.01499939, -0.007286072, 0.009933472, 0.06390381, 0.02444458, -0.010345459, 0.041931152, 0.032989502, -0.04522705, 0.056610107, 0.0068893433, -0.008911133, 0.012489319, 0.01675415, 0.020065308, 0.018753052, 0.022659302, -0.051849365, -0.04925537, 0.046325684, -0.005268097, 0.0026874542, -0.036712646, 0.009437561, -0.0037841797, -0.01473999, -0.034179688, -0.0011606216, 0.05026245, 0.0020771027, -0.016021729, -0.0044898987, 0.04168701, -0.015205383, 0.019210815, -0.012374878, -0.031311035, 0.03111267, -0.040100098, -0.016479492, 0.020446777, 0.010192871, 0.0037841797, -0.0023765564, 0.015220642, -0.016571045, -0.006454468, 0.037384033, -0.044555664, -0.008262634, 0.019546509, 0.009460449, 0.014701843, 0.02658081, -0.02078247, 0.015571594, 0.013153076, -0.010375977, 0.047912598, 0.005393982, -0.007911682, -0.019378662, 0.023529053, -0.0033550262, -0.04598999, -0.0052871704, 0.040252686, 0.011375427, 0.01550293, -0.004508972, 0.006515503, 0.003370285, -0.022766113, 0.00062561035, -0.0007596016, -0.0015277863, 0.0149002075, 0.061401367, 8.261204e-05, 0.06359863, -0.01537323, 0.007446289, 0.018814087, 0.02507019, 0.024215698, 0.006122589, 0.005886078, -0.03829956, 0.029037476, 0.07720947, 0.016921997, 0.022109985, 0.005958557, 0.028793335, 0.019485474, 0.015174866, 0.026153564, 0.032318115, 0.034210205, 0.027145386, -0.019515991, -0.018661499, 0.020477295, 0.008598328, -0.06573486, -0.037109375, 0.04043579, 0.030471802, -0.0010843277, 0.009757996, 0.026947021, 0.037017822, -0.018234253, -0.0115356445, 0.099365234, 0.027816772, -0.019927979, 0.0020961761, 0.013198853, -0.019073486, 2.7656555e-05, 0.041259766, 0.029510498, -0.016204834, 0.028137207, 0.039489746, 0.034698486, -0.03918457, -0.029418945, 0.02041626, 0.0073432922, -0.018569946, -0.009849548, 0.002861023, 0.030319214, -0.012886047, 0.014671326, -0.035827637, 0.007247925, -0.027709961, -0.022079468, 0.0012960434, 0.015426636, -0.01725769, 0.01525116, 0.025360107, -0.0077400208, -0.039916992, 0.029037476, -0.011154175, 0.007736206, -0.041748047, 0.05343628, 0.007286072, 0.0435791, 0.034301758, -0.047210693, 0.03552246, -0.015327454, 0.029922485, -0.018859863, 0.013053894, -0.028060913, 0.07757568, -0.020462036, 0.070739746, -0.010223389, 0.03604126, 0.02758789, -0.023284912, 0.012184143, 0.029144287, 0.023880005, -0.019378662, -0.0051116943, 0.0048675537, 0.01864624, -0.04397583, -0.007598877, 0.0713501, 0.0115737915, 0.002922058, 0.011619568, 0.017364502, 0.031921387, -0.0019664764, -0.008575439, 0.003484726, -0.09466553, 0.03475952, 0.026611328, -0.039520264, -0.0104522705, -0.005443573, -0.008392334, 0.012908936, 0.0043792725, -0.002456665, -0.028396606, -0.02027893, -0.0005569458, 0.027786255, 0.03427124, -0.0062332153, -0.018203735, 0.019241333, 0.07244873, -0.0028057098, 0.01234436, -0.0018787384, -0.027496338, 0.0015287399, -0.004032135, -0.013748169, -0.01878357, 0.0018053055, -0.01159668, 0.028213501, 0.004776001, 0.042388916, 0.0024280548, 0.017471313, -0.038085938, 0.026321411, 0.02973938, 0.06213379, 0.006401062, 0.036102295, -0.028121948, -0.00869751, -0.016693115, 0.029190063, 0.016784668, -0.008628845, 0.0039634705, -0.0035381317, 0.019500732, 0.025009155, -0.04547119, -0.003572464, 0.05215454, 0.067871094, -0.04257202, -0.02293396, -0.027175903, 0.05340576, 0.019226074, 0.039978027, 0.056121826, -0.028320312, -0.020217896, -0.035003662, 0.03225708, 0.028656006, 0.062347412, 0.12915039, -0.0137786865, 0.0022201538, -0.057434082, -0.04397583, -0.049865723, -0.013160706, -0.03353882, 0.006427765, -0.014823914, -0.008201599, -0.036346436, -0.037353516, -0.010528564, -0.015930176, -0.027572632, 0.0074272156, 0.004547119, -0.024414062, -0.018859863, -0.020095825, 0.029632568, -0.00067043304, -0.044036865, -0.0043411255, -0.005256653, -0.019195557, 0.022262573, -0.00020956993, -0.013877869, -0.011108398, -0.020324707, -0.015808105, -0.025039673, -0.009498596, 0.05090332, 0.0046195984, -0.017150879, 0.04309082, -0.029067993, 0.002670288, -0.00026249886, -0.032409668, -0.053100586, 0.012481689, -0.014633179, 0.0013475418, -0.034332275, 0.038330078, 0.014892578, -0.046936035, 0.021591187, -0.020385742, -0.0052604675, 0.02796936, 0.0014333725, 0.012077332, -0.0118255615, -0.005569458, 0.008491516, 0.009841919, 0.0031318665, -0.003408432, -0.007144928, 0.040374756, -0.0038928986, 0.005279541, -0.008415222, 0.031707764, 0.0140686035, -0.015029907, -0.02810669, -0.0078125, -0.030853271, -0.03201294, 0.021316528, -0.036193848, -0.0423584, 0.0072784424, 0.014801025, 0.0019607544, -0.012367249, -0.009056091, -0.021438599, -0.02645874, 0.038726807, -0.007549286, 0.0049591064, 0.019012451, 0.017791748, -0.009185791, 0.04006958, 0.003107071, -0.0075302124, -0.010375977, -0.009246826, -0.02130127, -0.0056762695, -0.0076789856, 0.010009766, -0.010536194, 0.041107178, 0.0021133423, 0.029891968, 0.01626587, 0.042236328, -0.02784729, -0.032836914, 0.0317688, 0.045715332, 0.000116825104, 0.028030396, 0.007205963, 0.012512207, -0.035583496, -0.048034668, -0.023529053, -0.04953003, 0.0345459, -0.048339844, -0.060272217, -0.004512787, 0.04425049, 0.0076141357, 0.029510498, 0.007396698, 0.003353119, -0.038726807, 0.07183838, -0.026901245, -0.023529053, -0.038085938, 0.068725586, 0.018096924, -0.013534546, 0.05883789, -0.016113281, 0.017944336, 0.041046143, 0.022918701, 0.036499023, 0.015296936, -0.04916382, 0.0075683594, -0.011390686, 0.009735107, -0.0070152283, 0.003129959, -0.032562256, 0.0003478527, -0.0036640167, -0.006893158, -0.016098022, -0.034332275, 0.037750244, -0.010269165, 0.016494751, -0.02394104, 0.03753662, -0.022644043, -0.0008234978, 0.001001358, -0.048217773, 0.04989624, 0.0078125, 0.0044937134, 0.027038574, 0.04736328, -0.02973938, -0.011726379, 0.01348114, 0.021408081, 0.00844574, -0.03741455, -0.015686035, -0.040893555, 0.001452446, -0.025405884, 0.07348633, 0.038238525, -0.019958496, 0.023071289, -0.016403198, -0.08105469, 0.0071029663, -0.019088745, 5.8174133e-05, -0.005569458, 0.01399231, 0.02255249, 0.011222839, 0.00028824806, 0.0066184998, 0.0017499924, -0.009864807, -0.0115737915, 0.053100586, 0.0065231323, 0.001865387, -0.026428223, 0.03692627, 0.025390625, 0.022613525, 0.018722534, 0.007675171, -0.03439331, 0.041625977, -0.01789856, -0.041046143, 0.0051460266, 0.04144287, 0.048553467, 0.054595947, -0.01108551, -0.033935547, -0.026275635, -0.0118255615, -0.021362305, -0.009841919, -0.00724411, 0.028900146, 0.009887695, -0.023803711, 0.016311646, 0.018798828, -0.03668213, 0.046844482, 0.010696411, -0.014717102, -0.008110046, -0.004589081, -0.0028076172, -0.050811768, -0.017196655, -0.03491211, 0.0074005127, -0.038909912, 0.032440186, -0.034362793, -0.008682251, 0.032928467, -0.04626465, -0.009666443, 0.018951416, 0.031951904, -0.003791809, 0.02015686, -0.05532837, -0.005683899, -0.00054216385, -0.0034332275, 0.008659363, 0.02130127, -0.038879395, -0.0033397675, -0.03866577, -0.0049934387, 0.017944336, 0.001496315, 0.019485474, -0.004348755, 0.00046491623, 0.0007157326, 0.035614014, -0.027694702, 0.03692627, -0.008491516, 0.0524292, -0.016662598, -0.0017795563, -0.021575928, -0.018753052, -0.049346924, -0.06652832, 0.04272461, 0.03186035, 0.0011978149, 0.03463745, 0.024002075, 0.02607727, 0.020446777, 0.0256958, 0.026855469, 0.0074005127, -0.067993164, 0.017944336, -0.0039482117, 0.05496216, -0.041412354, 0.014175415, 0.02444458, -0.026412964, 0.057403564, -0.026779175, 0.023254395, 0.03945923, 0.033569336, -0.030258179, -0.039093018, -0.036468506, 0.017105103, 0.009635925, 0.025497437, 0.04156494, -0.02571106, -0.0010414124, -0.005630493, -0.016448975, -0.026733398, 0.001326561, -0.042022705, 0.0012521744, -0.041259766, -0.12182617, -0.03857422, 0.12548828, -0.005947113, -0.020736694, -0.0033855438, 0.03778076, -0.033813477, 0.038970947, 0.003921509, 0.011810303, 0.031982422, -0.032562256, -0.002653122, -0.025009155, -0.03805542, -0.016998291, 0.018173218, 0.0158844, 0.0011739731, 0.048217773, -0.020401001, 0.044708252, -0.017318726, 0.014457703, -0.041809082, 0.010543823, 0.041931152, 0.076293945, -0.054779053, 0.060272217, -0.046936035, 0.02949524, 0.00554657, 0.041534424, -0.013046265, -0.056152344, 0.010406494, 0.02973938, -0.023727417, -0.022476196, -0.024734497, -0.013168335, 0.060424805, 0.011787415, 0.018997192, -0.043426514, -0.00077724457, -0.010154724, 0.017150879, -0.01171875, -0.022476196, 0.0034255981, -0.0026454926, 0.004837036, -0.0043296814, 0.02619934, -0.021560669, -0.039733887, -0.022415161, -0.06817627, -0.023223877, -0.018585205, -0.015319824, 0.012588501, 0.0064353943, -0.013748169, 0.043304443, 0.002626419, -0.029373169, -0.016784668, -0.026184082, 0.05847168, 0.034179688, 0.03842163, -0.05493164, -0.017486572, 0.016540527, 0.03164673, 0.089904785, 0.013534546, -0.07684326, -0.024108887, 0.07434082, 0.030395508, 0.007091522, 0.07373047, 0.012527466, -0.010856628, -0.01828003, -0.045196533, 0.00065279007, -0.0637207, 0.010726929, 0.023880005, -0.0030708313, -0.012298584, 0.027236938, -0.04928589, 0.023071289, 0.008674622, -0.023529053, -0.015838623, -0.010543823, 0.012168884, 0.014854431, -0.05834961, -0.06088257, -0.012313843, 0.035461426, 0.02027893, 0.019348145, -0.014602661, -0.02104187, -0.0309906, 0.001405716, -0.019973755, -0.00157547, -0.003944397, 0.0009326935, -0.02078247, -0.015731812, -0.044433594, 0.03390503, 0.057159424, 0.018585205, -0.023895264, -0.0057029724, 0.0049552917, 0.013412476, 0.022399902, 0.010154724, 0.0519104, 0.06591797, 0.018341064, 0.012161255, -0.05810547, -0.043304443, -0.031173706, 0.0023860931, -0.003944397, 0.11425781, -0.031036377, 0.019989014, -0.038635254, -0.025939941, 0.035064697, 0.041168213, 0.03161621, -0.069885254, -0.04537964, 0.028945923, -0.023162842, 0.019226074, -0.028442383, 0.015594482, -0.019256592, -0.0046463013, 0.034240723, 0.009124756, 0.05718994, 0.031219482, 0.02154541, 0.009590149, 0.00076818466, 0.04849243, -0.029129028, -0.03375244, -0.023391724, -0.028381348, -0.029708862, -0.0132369995, 0.010353088, 0.020263672, -0.030807495, 0.01007843, -0.03704834, 0.023376465, -0.03665161, 0.03741455, 0.015144348, 0.057281494, 0.03137207, 0.048431396, 0.021194458, 0.008110046, -0.03540039, -0.015312195, 0.022384644, 0.0065956116, 0.008056641, 0.0018348694, -0.009246826, 0.030380249, 0.0003862381, 0.0051841736, 0.04486084, 0.017807007, 0.0026130676, 0.07977295, 0.05419922, 0.062194824, 0.02633667, 0.024841309, -0.041625977, -0.005897522, 0.04031372, -0.055908203, 0.0026226044, -0.05340576, -0.05496216, 0.011474609, -0.006954193, -0.013122559, 0.019714355, -0.07159424, 0.031173706, 0.0034255981, -0.0034103394, 0.0440979, 0.011779785, -0.007827759, -0.03173828, -0.020950317, -0.030166626, -0.035308838, 0.030792236, 0.04525757, -0.028701782, -0.011100769, -0.02331543, -0.0357666, -0.025680542, 0.0011911392, 0.01940918, 0.05706787, 0.028381348, 0.007133484, -0.07733154, -0.007686615, 0.03869629, 0.0066833496, 0.008842468, 0.03439331, -0.014282227, 0.0357666, -0.004737854, -0.039794922, -0.0070381165, 0.02670288, 0.0107421875, 0.016189575, -0.06555176, -0.0138549805, 0.0008363724, -0.016693115, 0.006904602, -0.020263672, -0.030426025, 0.008453369, -0.046173096, -0.01802063, -0.013595581, -0.0044288635, -0.0039978027, -0.0044898987, 0.0007619858, 0.003921509, 0.0053977966, 0.020385742, -0.012329102, -0.023803711, -0.0057525635, 0.038330078, -0.014549255, -0.06298828, -0.047607422, 0.039245605, -0.06781006, -0.035217285, -0.009056091, 0.019927979, -0.003932953, -0.020309448, -0.017044067, 0.018127441, -8.624792e-05, -0.043182373, 0.009590149, 0.035308838, 0.031951904, 0.0011615753, -0.042022705, 0.079956055, 0.026687622, 0.013542175, -0.0074157715, -0.00983429, -0.0022563934, 0.07373047, 0.059387207, 0.03488159, 0.0071372986, -0.06427002, -0.0546875, -0.02482605, 0.11071777, -0.021072388, 0.01626587, -0.049713135, 0.061553955, -0.016860962, 0.051971436, -0.012962341, -0.0011711121, -0.014198303, -0.0061149597, -0.005836487, 0.00022387505, -0.027618408, 0.019836426, 0.009933472, 0.02368164, -0.020309448, -0.0049591064, -0.008628845, -0.03253174, -0.017684937, 0.02468872, -0.0023498535, 0.01448822, 0.061920166, 0.031707764, -0.0026416779, -0.040985107, -0.06335449, -0.036071777, 0.05404663, -0.0044136047, -0.0146102905, -0.0033416748, 0.028671265, -0.012771606, -0.0016565323, -0.0038909912, -0.02407837, -0.009857178, 0.0014467239, -0.008720398, -0.006011963, 0.032073975, -0.033325195, 0.014862061, -0.017227173, -0.018753052, -0.0060424805, 0.022567749, -0.017654419, -0.017562866, -0.07244873, -0.0881958, 0.050476074, 0.02609253, -0.032409668, 0.07458496, 0.009399414, 0.009117126, -0.031051636, -0.03451538, -0.004219055, -0.05718994, 0.020080566, -0.025421143, -0.010948181, 0.06341553, -0.009231567, -0.021697998, -0.009719849, 0.012802124, -0.020370483, 0.0034389496, 0.018859863, -0.025680542, 0.0013141632, 0.068603516, -0.021026611, 0.021881104, -0.0395813, -0.0019073486, 0.0056037903, -0.032348633]

Retrieve the most relevant chunks from the vector database

We use cosine similarity to find the most similar chunks

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

similarities = [cosine_similarity(query_embedding, chunk) for chunk in embeddings]
print("similarity scores: ", similarities)

sorted_indices = np.argsort(similarities)[::-1]

top_indices = sorted_indices[:10]
print("Here are the indices of the top 10 chunks after retrieval: ", top_indices)

top_chunks_after_retrieval = [chunks[i] for i in top_indices]
print("Here are the top 10 chunks after retrieval: ")
for t in top_chunks_after_retrieval:
    print("== " + t)
similarity scores:  [0.6953257882233425, 0.3713410510180273, 0.46501499776898964, 0.5448546916785195, 0.4014738351361969, 0.3231420292334584, 0.3179003053384008, 0.42799691553367775, 0.18882594531435784, 0.36868801306504106, 0.3404040737300553, 0.3852837621219358, 0.2600249419491577, 0.3723244353775111, 0.3631492691137214, 0.47574774051439606, 0.40415422750911745, 0.4149923346201023, 0.5014741934381444, 0.3549433331883204, 0.32072714802512714, 0.14770872479410424, 0.585277816615252, 0.6999636953772764, 0.7722295084104617, 0.4895347049465806, 0.5170096485954725, 0.7137817366881455, 0.5224900699612323, 0.5914632581598285, 0.2657897083381463, 0.6462342489537262, 0.6317222315431096, 0.5982303530756702, 0.5138265091630297, 0.41385121172723643, 0.4293941094100836, 0.4173182546482015, 0.42621236706314475, 0.4428474375355954, 0.35058541576139896, 0.3578709652019502, 0.3930157841938308, 0.3564608202848675, 0.23016661533167404, 0.4933441863421645, 0.41037089239250985, 0.39993051898770193, 0.3119997063424595, 0.2677143729521374, 0.3700866951454496, 0.46727994925061545, 0.4393343280374425, 0.42111290117172434, 0.4485349189824285, 0.4710573736688592, 0.24169956903740436, 0.3840442910806355, 0.14284631817675886, 0.5381588054138154, 0.431113882725076, 0.5189547209048608, 0.3950667224233914, 0.32429768756510174, 0.4370358125161736, 0.18727062244331039, 0.5206375682478743, 0.5175737635701252, 0.5326043981628349, 0.45586923626994363, 0.21667338125532032, 0.16459878595959285, 0.22236726481673777, 0.5187259906958807, 0.2884444442338396, 0.286407544555338, 0.2313840178160818, 0.2057731158935257, 0.5973876998341746, 0.42904243401792086, 0.4081217538000544, 0.5330523063972133, 0.45080561486977405, 0.414703452285757, 0.2569028899107211, 0.5087916806929323, 0.14159076456040554, 0.46505779053352697, 0.556364222182839, 0.35464351181035236, 0.40174477023626]
Here are the indices of the top 10 chunks after retrieval:  [24 27 23  0 31 32 33 78 29 22]
Here are the top 10 chunks after retrieval: 
== stunt coordinator. Dune: Part Two was produced by Villeneuve, Mary Parent, and Cale Boyter, with Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Jon Spaihts, Richard P. Rubinstein, John Harrison, and Herbert W. Gain serving as executive producers and Kevin J. Anderson as creative consultant. Legendary CEO Joshua Grode confirmed in April 2019 that they plan to make a sequel, adding that "there's a logical place to stop the [first] movie before the book is over".In December 2020,
== that."On October 26, 2021, Legendary officially greenlit Dune: Part Two, with a spokesperson for the company stating, "We would not have gotten to this point without the extraordinary vision of Denis and the amazing work of his talented crew, the writers, our stellar cast, our partners at Warner Bros., and of course the fans! Here's to more Dune." Production work had occurred back-to-back with the first film, as Villeneuve and his wife Lapointe immediately took a flight to Budapest in order to begin
== series. Villeneuve ultimately secured a two-film deal with Warner Bros. Pictures, in the same style as the two-part adaption of Stephen King's It in 2017 and 2019. In January 2019, Joe Walker was confirmed to be serving as the film's editor. Other crew included Brad Riker as supervising art director, Patrice Vermette as production designer, Paul Lambert as visual effects supervisor, Gerd Nefzer as special effects supervisor, and Thomas Struthers as stunt coordinator. Dune: Part Two was produced by
== Dune: Part Two is a 2024 American epic science fiction film directed and produced by Denis Villeneuve, who co-wrote the screenplay with Jon Spaihts. The sequel to Dune (2021) adapts the 1965 novel Dune by Frank Herbert and follows Paul Atreides as he unites with the Fremen people of the desert planet Arrakis to wage war against House Harkonnen. TimothΓ©e Chalamet, Rebecca Ferguson, Josh Brolin, Stellan SkarsgΓ₯rd, Dave Bautista, Zendaya, Charlotte Rampling, and Javier Bardem reprise their roles from the first
== Eric Roth was hired to co-write the screenplay in April 2017 for the Dune films, and Jon Spaihts was later confirmed to be co-writing the script alongside Roth and Villeneuve. Game of Thrones language creator David Peterson was confirmed to be developing languages for the film in April 2019. Villeneuve and Peterson had created the Chakobsa language, which was used by actors on set. In November 2019, Spaihts stepped down as showrunner for Dune: Prophecy to focus on Dune: Part Two. In June 2020, Greig Fraser
== on Dune: Part Two. In June 2020, Greig Fraser said, "It's a fully formed story in itself with places to go. It's a fully standalone epic film that people will get a lot out of when they see it". Between the release of Dune and the confirmation of Dune: Part Two, Villeneuve started working the script in a way that production could begin immediately once the film was greenlit. By February 2021, Roth created a full treatment for the sequel, with writing beginning that August. He confirmed that Feyd-Rautha
== that August. He confirmed that Feyd-Rautha would appear in the film, and stated he will be a "very important character". In March 2022, Villeneuve had mostly finished writing the screenplay. Craig Mazin and Roth wrote additional literary material for the film.Villeneuve stated that the film would continue directly from the first, and specifically described it as being the "second part." He described the film as being an "epic war movie", adding that while the first film was more "contemplative", the second
== On the review aggregator website Rotten Tomatoes, 93% of 378 critics' reviews are positive, with an average rating of 8.4/10. The website's consensus reads: "Visually thrilling and narratively epic, Dune: Part Two continues Denis Villeneuve's adaptation of the beloved sci-fi series in spectacular form." Metacritic, which uses a weighted average, assigned the film a score of 79 out of 100, based on 62 critics, indicating "generally favorable" reviews. Audiences surveyed by CinemaScore gave the film an
== theatrical experience is at the very heart of the cinematic language for me." With Dune: Part Two being greenlit, Villeneuve said that his primary concern was to complete the filming as soon as possible, with the earliest he expected to start in the last quarter of 2022. However, he noted that production would be facilitated by the work already established on the first film, which would help expedite production.
== By November 2016, Legendary Pictures had obtained the film and TV rights for the Dune franchise, based on the eponymous 1965 novel by Frank Herbert. Vice chair of worldwide production for Legendary Mary Parent began discussing with Denis Villeneuve about directing a film adaptation, quickly hiring him after realizing his passion for Dune. By February 2018, Villeneuve was confirmed to be hired as director, and intended to adapt the novel as a two-part film series. Villeneuve ultimately secured a two-film

Step 2 - Rerank the chunks retrieved from the vector database

We rerank the 10 chunks retrieved from the vector database. Reranking boosts retrieval accuracy.

Reranking lets us go from 10 chunks retrieved from the vector database, to the 3 most relevant chunks.

response = co.rerank(
    query=query,
    documents=top_chunks_after_retrieval,
    top_n=3,
    model="rerank-english-v2.0",
)

top_chunks_after_rerank = [result.document['text'] for result in response]
print("Here are the top 3 chunks after rerank: ")
for t in top_chunks_after_rerank:
    print("== " + t)
Here are the top 3 chunks after rerank: 
== Dune: Part Two is a 2024 American epic science fiction film directed and produced by Denis Villeneuve, who co-wrote the screenplay with Jon Spaihts. The sequel to Dune (2021) adapts the 1965 novel Dune by Frank Herbert and follows Paul Atreides as he unites with the Fremen people of the desert planet Arrakis to wage war against House Harkonnen. TimothΓ©e Chalamet, Rebecca Ferguson, Josh Brolin, Stellan SkarsgΓ₯rd, Dave Bautista, Zendaya, Charlotte Rampling, and Javier Bardem reprise their roles from the first
== stunt coordinator. Dune: Part Two was produced by Villeneuve, Mary Parent, and Cale Boyter, with Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Jon Spaihts, Richard P. Rubinstein, John Harrison, and Herbert W. Gain serving as executive producers and Kevin J. Anderson as creative consultant. Legendary CEO Joshua Grode confirmed in April 2019 that they plan to make a sequel, adding that "there's a logical place to stop the [first] movie before the book is over".In December 2020,
== series. Villeneuve ultimately secured a two-film deal with Warner Bros. Pictures, in the same style as the two-part adaption of Stephen King's It in 2017 and 2019. In January 2019, Joe Walker was confirmed to be serving as the film's editor. Other crew included Brad Riker as supervising art director, Patrice Vermette as production designer, Paul Lambert as visual effects supervisor, Gerd Nefzer as special effects supervisor, and Thomas Struthers as stunt coordinator. Dune: Part Two was produced by

Step 3 - Generate the model final answer, given the retrieved and reranked chunks

preamble = """
## Task & Context
You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.

## Style Guide
Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.
"""
documents = [
    {"title": "chunk 0", "snippet": top_chunks_after_rerank[0]},
    {"title": "chunk 1", "snippet": top_chunks_after_rerank[1]},
    {"title": "chunk 2", "snippet": top_chunks_after_rerank[2]},
  ]

response = co.chat(
  message=query,
  documents=documents,
  preamble=preamble,
  model="command-r",
  temperature=0.3
)

print("Final answer:")
print(response.text)
Final answer:
Here are the key crew members involved in 'Dune: Part Two':

- Denis Villeneuve: director and producer
- Jon Spaihts: co-writer of the screenplay
- Mary Parent and Cale Boyter: producers 
- Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Richard P. Rubinstein, John Harrison, Herbert W. Gain and Kevin J. Anderson: executive producers 
- Joe Walker: editor
- Brad Riker: supervising art director
- Patrice Vermette: production designer
- Paul Lambert: visual effects supervisor
- Gerd Nefzer: special effects supervisor
- Thomas Struthers: stunt coordinator. 

A number of crew members from the first film returned for the sequel, which is set to be released in 2024.

Note: this is indeed the answer you'd expect, and here was the passage of text in wikipedia explaining it!

" [...] Dune: Part Two was originally scheduled to be released on October 20, 2023, but was delayed to November 17, 2023, before moving forward two weeks to November 3, 2023, to adjust to changes in release schedules from other studios. It was later postponed by over four months to March 15, 2024, due to the 2023 Hollywood labor disputes. After the strikes were resolved, the film moved once more up two weeks to March 1, 2024. [...]"

Bonus: Citations come for free with Cohere! πŸŽ‰

At Cohere, all RAG calls come with... precise citations! πŸŽ‰
The model cites which groups of words, in the RAG chunks, were used to generate the final answer.
These citations make it easy to check where the model’s generated response claims are coming from.
They help users gain visibility into the model reasoning, and sanity check the final model generation.
These citations are optional β€” you can decide to ignore them.

print("Citations that support the final answer:")
for cite in response.citations:
  print(cite)
Citations that support the final answer:
{'start': 63, 'end': 79, 'text': 'Denis Villeneuve', 'document_ids': ['doc_0']}
{'start': 81, 'end': 102, 'text': 'director and producer', 'document_ids': ['doc_0']}
{'start': 105, 'end': 116, 'text': 'Jon Spaihts', 'document_ids': ['doc_0']}
{'start': 118, 'end': 145, 'text': 'co-writer of the screenplay', 'document_ids': ['doc_0']}
{'start': 148, 'end': 159, 'text': 'Mary Parent', 'document_ids': ['doc_1']}
{'start': 164, 'end': 175, 'text': 'Cale Boyter', 'document_ids': ['doc_1']}
{'start': 177, 'end': 186, 'text': 'producers', 'document_ids': ['doc_1']}
{'start': 190, 'end': 204, 'text': 'Tanya Lapointe', 'document_ids': ['doc_1']}
{'start': 206, 'end': 219, 'text': 'Brian Herbert', 'document_ids': ['doc_1']}
{'start': 221, 'end': 234, 'text': 'Byron Merritt', 'document_ids': ['doc_1']}
{'start': 236, 'end': 247, 'text': 'Kim Herbert', 'document_ids': ['doc_1']}
{'start': 249, 'end': 260, 'text': 'Thomas Tull', 'document_ids': ['doc_1']}
{'start': 262, 'end': 283, 'text': 'Richard P. Rubinstein', 'document_ids': ['doc_1']}
{'start': 285, 'end': 298, 'text': 'John Harrison', 'document_ids': ['doc_1']}
{'start': 300, 'end': 315, 'text': 'Herbert W. Gain', 'document_ids': ['doc_1']}
{'start': 320, 'end': 337, 'text': 'Kevin J. Anderson', 'document_ids': ['doc_1']}
{'start': 339, 'end': 358, 'text': 'executive producers', 'document_ids': ['doc_1']}
{'start': 362, 'end': 372, 'text': 'Joe Walker', 'document_ids': ['doc_2']}
{'start': 374, 'end': 380, 'text': 'editor', 'document_ids': ['doc_2']}
{'start': 383, 'end': 393, 'text': 'Brad Riker', 'document_ids': ['doc_2']}
{'start': 395, 'end': 419, 'text': 'supervising art director', 'document_ids': ['doc_2']}
{'start': 422, 'end': 438, 'text': 'Patrice Vermette', 'document_ids': ['doc_2']}
{'start': 440, 'end': 459, 'text': 'production designer', 'document_ids': ['doc_2']}
{'start': 462, 'end': 474, 'text': 'Paul Lambert', 'document_ids': ['doc_2']}
{'start': 476, 'end': 501, 'text': 'visual effects supervisor', 'document_ids': ['doc_2']}
{'start': 504, 'end': 515, 'text': 'Gerd Nefzer', 'document_ids': ['doc_2']}
{'start': 517, 'end': 543, 'text': 'special effects supervisor', 'document_ids': ['doc_2']}
{'start': 546, 'end': 562, 'text': 'Thomas Struthers', 'document_ids': ['doc_2']}
{'start': 564, 'end': 582, 'text': 'stunt coordinator.', 'document_ids': ['doc_2']}
{'start': 686, 'end': 691, 'text': '2024.', 'document_ids': ['doc_0']}
def insert_citations_in_order(text, citations):
    """
    A helper function to pretty print citations.
    """
    offset = 0
    document_id_to_number = {}
    citation_number = 0
    modified_citations = []

    # Process citations, assigning numbers based on unique document_ids
    for citation in citations:
        citation_numbers = []
        for document_id in sorted(citation["document_ids"]):
            if document_id not in document_id_to_number:
                citation_number += 1  # Increment for a new document_id
                document_id_to_number[document_id] = citation_number
            citation_numbers.append(document_id_to_number[document_id])

        # Adjust start/end with offset
        start, end = citation['start'] + offset, citation['end'] + offset
        placeholder = ''.join([f'[{number}]' for number in citation_numbers])
        # Bold the cited text and append the placeholder
        modification = f'**{text[start:end]}**{placeholder}'
        # Replace the cited text with its bolded version + placeholder
        text = text[:start] + modification + text[end:]
        # Update the offset for subsequent replacements
        offset += len(modification) - (end - start)

    # Prepare citations for listing at the bottom, ensuring unique document_ids are listed once
    unique_citations = {number: doc_id for doc_id, number in document_id_to_number.items()}
    citation_list = '\n'.join([f'[{doc_id}] source: "{documents[doc_id - 1]["snippet"]}"' for doc_id, number in sorted(unique_citations.items(), key=lambda item: item[1])])
    text_with_citations = f'{text}\n\n{citation_list}'

    return text_with_citations


print(insert_citations_in_order(response.text, response.citations))

Here are the key crew members involved in 'Dune: Part Two':

- **Denis Villeneuve**[1]: **director and producer**[1]
- **Jon Spaihts**[1]: **co-writer of the screenplay**[1]
- **Mary Parent**[2] and **Cale Boyter**[2]: **producers**[2] 
- **Tanya Lapointe**[2], **Brian Herbert**[2], **Byron Merritt**[2], **Kim Herbert**[2], **Thomas Tull**[2], **Richard P. Rubinstein**[2], **John Harrison**[2], **Herbert W. Gain**[2] and **Kevin J. Anderson**[2]: **executive producers**[2] 
- **Joe Walker**[3]: **editor**[3]
- **Brad Riker**[3]: **supervising art director**[3]
- **Patrice Vermette**[3]: **production designer**[3]
- **Paul Lambert**[3]: **visual effects supervisor**[3]
- **Gerd Nefzer**[3]: **special effects supervisor**[3]
- **Thomas Struthers**[3]: **stunt coordinator.**[3] 

A number of crew members from the first film returned for the sequel, which is set to be released in **2024.**[1]

[1] source: "Dune: Part Two is a 2024 American epic science fiction film directed and produced by Denis Villeneuve, who co-wrote the screenplay with Jon Spaihts. The sequel to Dune (2021) adapts the 1965 novel Dune by Frank Herbert and follows Paul Atreides as he unites with the Fremen people of the desert planet Arrakis to wage war against House Harkonnen. TimothΓ©e Chalamet, Rebecca Ferguson, Josh Brolin, Stellan SkarsgΓ₯rd, Dave Bautista, Zendaya, Charlotte Rampling, and Javier Bardem reprise their roles from the first"
[2] source: "stunt coordinator. Dune: Part Two was produced by Villeneuve, Mary Parent, and Cale Boyter, with Tanya Lapointe, Brian Herbert, Byron Merritt, Kim Herbert, Thomas Tull, Jon Spaihts, Richard P. Rubinstein, John Harrison, and Herbert W. Gain serving as executive producers and Kevin J. Anderson as creative consultant. Legendary CEO Joshua Grode confirmed in April 2019 that they plan to make a sequel, adding that "there's a logical place to stop the [first] movie before the book is over".In December 2020,"
[3] source: "series. Villeneuve ultimately secured a two-film deal with Warner Bros. Pictures, in the same style as the two-part adaption of Stephen King's It in 2017 and 2019. In January 2019, Joe Walker was confirmed to be serving as the film's editor. Other crew included Brad Riker as supervising art director, Patrice Vermette as production designer, Paul Lambert as visual effects supervisor, Gerd Nefzer as special effects supervisor, and Thomas Struthers as stunt coordinator. Dune: Part Two was produced by"