LlamaIndex

Prerequisite

To use LlamaIndex and Cohere, you will need:

LlamaIndex Package. To install it, run pip install llama-index.
Cohere’s SDK. To install it, run pip install cohere. If you run into any issues or want more details on Cohere’s SDK, see this wiki.
A Cohere API Key. For more details on pricing see this page. When you create an account with Cohere, we automatically create a trial API key for you. This key will be available on the dashboard where you can copy it, and it’s in the dashboard section called “API Keys” as well.

Cohere Chat with LlamaIndex

To use Cohere’s chat functionality with LlamaIndex create a Cohere model object and call the chat function.

PYTHON

1 from llama_index.llms.cohere import Cohere
2 from llama_index.core.llms import ChatMessage
3 
4 cohere_model = Cohere(api_key="{API_KEY}")
5 message = ChatMessage(role="user",content= "Who founded Cohere?")
6 resp = cohere_model.chat([message])
7 print(resp)

Cohere Embeddings with LlamaIndex

To use Cohere’s embeddings with LlamaIndex create a Cohere Embeddings object with an embedding model from this list and call get_text_embedding.

PYTHON

1 from llama_index.embeddings.cohereai import CohereEmbedding
2 
3 embed_model = CohereEmbedding(
4     cohere_api_key="{API_KEY}",
5     model_name="embed-english-v3.0", # Supports all Cohere embed models
6     input_type="search_query", # Required for v3 models
7 )
8 
9 # Generate Embeddings
10 embeddings = embed_model.get_text_embedding("Welcome to Cohere!")
11 
12 # Print embeddings
13 print(len(embeddings))
14 print(embeddings[:5])

Cohere Rerank with LlamaIndex

To use Cohere’s rerank functionality with LlamaIndex create a Cohere Rerank object and use as a node post processor.

PYTHON

1 cohere_rerank = CohereRerank(api_key="{API_KEY}", top_n=2)

Cohere Pipeline with LlamaIndex

The following example uses Cohere’s chat model, embeddings and rerank functionality to generate a response based on your data.

PYTHON

1 # rerank
2 
3 from llama_index import ServiceContext, VectorStoreIndex
4 from llama_index.llms.cohere import Cohere
5 from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
6 from llama_index.embeddings.cohereai import CohereEmbedding
7 from llama_index.postprocessor.cohere_rerank import CohereRerank
8 
9 # Create the embedding model
10 embed_model = CohereEmbedding(
11     cohere_api_key="{API_KEY}",
12     model_name="embed-english-v3.0",
13     input_type="search_query",
14 )
15 
16 # Create the service context with the cohere model for generation and embedding model
17 service_context = ServiceContext.from_defaults(
18     llm=Cohere(api_key="{API_KEY}", model="command"),
19     embed_model=embed_model
20 )
21 
22 # Load the data, for this example data needs to be in a test file 
23 data = SimpleDirectoryReader(input_files=["example_data_file.txt"]).load_data()
24 index = VectorStoreIndex.from_documents(data, service_context=service_context)
25 
26 # Create a cohere reranker 
27 cohere_rerank = CohereRerank(api_key="{API_KEY}")
28 
29 # Create the query engine
30 query_engine = index.as_query_engine(node_postprocessors=[cohere_rerank])
31 
32 # Generate the response
33 response = query_engine.query("Who founded Cohere?",)

1	from llama_index.llms.cohere import Cohere
2	from llama_index.core.llms import ChatMessage
3
4	cohere_model = Cohere(api_key="{API_KEY}")
5	message = ChatMessage(role="user",content= "Who founded Cohere?")
6	resp = cohere_model.chat([message])
7	print(resp)

1	from llama_index.embeddings.cohereai import CohereEmbedding
2
3	embed_model = CohereEmbedding(
4	cohere_api_key="{API_KEY}",
5	model_name="embed-english-v3.0", # Supports all Cohere embed models
6	input_type="search_query", # Required for v3 models
7	)
8
9	# Generate Embeddings
10	embeddings = embed_model.get_text_embedding("Welcome to Cohere!")
11
12	# Print embeddings
13	print(len(embeddings))
14	print(embeddings[:5])

1	# rerank
2
3	from llama_index import ServiceContext, VectorStoreIndex
4	from llama_index.llms.cohere import Cohere
5	from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
6	from llama_index.embeddings.cohereai import CohereEmbedding
7	from llama_index.postprocessor.cohere_rerank import CohereRerank
8
9	# Create the embedding model
10	embed_model = CohereEmbedding(
11	cohere_api_key="{API_KEY}",
12	model_name="embed-english-v3.0",
13	input_type="search_query",
14	)
15
16	# Create the service context with the cohere model for generation and embedding model
17	service_context = ServiceContext.from_defaults(
18	llm=Cohere(api_key="{API_KEY}", model="command"),
19	embed_model=embed_model
20	)
21
22	# Load the data, for this example data needs to be in a test file
23	data = SimpleDirectoryReader(input_files=["example_data_file.txt"]).load_data()
24	index = VectorStoreIndex.from_documents(data, service_context=service_context)
25
26	# Create a cohere reranker
27	cohere_rerank = CohereRerank(api_key="{API_KEY}")
28
29	# Create the query engine
30	query_engine = index.as_query_engine(node_postprocessors=[cohere_rerank])
31
32	# Generate the response
33	response = query_engine.query("Who founded Cohere?",)