pip install cohere

import cohere 
# 替换自己的api keyco = cohere.Client("xxx") 
docs = ["Carson City is the capital city of the American state of Nevada.", "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.", "Capitalization or capitalisation in English grammar is the use of a capital letter at the start of a word. English usage varies from capitalization in other languages.", "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.", "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states."] 
response = co.rerank( model="rerank-english-v3.0", query="What is the capital of the United States?", documents=docs, top_n=len(docs)) print(response)

id='ca4ddc53-c16a-442c-bd94-fd3409cbc209' results=[RerankResponseResultsItem(document=RerankResponseResultsItemDocument(text='Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.'), index=3, relevance_score=0.999071), RerankResponseResultsItem(document=RerankResponseResultsItemDocument(text='Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.'), index=4, relevance_score=0.7867867), RerankResponseResultsItem(document=RerankResponseResultsItemDocument(text='Carson City is the capital city of the American state of Nevada.'), index=0, relevance_score=0.32713068), RerankResponseResultsItem(document=RerankResponseResultsItemDocument(text='The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.'), index=1, relevance_score=0.014899643), RerankResponseResultsItem(document=RerankResponseResultsItemDocument(text='Capitalization or capitalisation in English grammar is the use of a capital letter at the start of a word. English usage varies from capitalization in other languages.'), index=2, relevance_score=0.0029235429)] meta=ApiMeta(api_version=ApiMetaApiVersion(version='1', is_deprecated=None, is_experimental=None), billed_units=ApiMetaBilledUnits(input_tokens=None, output_tokens=None, search_units=1, classifications=None), tokens=None, warnings=None)

from langchain.retrievers.contextual_compression import ContextualCompressionRetrieverfrom langchain.retrievers.document_compressors import CohereRerankfrom langchain.llms.openai import OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitterfrom langchain_community.document_loaders import WebBaseLoaderfrom langchain_community.vectorstores import Chromafrom langchain_openai import OpenAIEmbeddings
urls = ["https://lilianweng.github.io/posts/2023-06-23-agent/",]
docs = [WebBaseLoader(url).load() for url in urls]docs_list = [item for sublist in docs for item in sublist]
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=250, chunk_overlap=0)doc_splits = text_splitter.split_documents(docs_list)
# Add to vectorDBvectorstore = Chroma.from_documents(documents=doc_splits,collection_name="rag-chroma",embedding=OpenAIEmbeddings(),)retriever = vectorstore.as_retriever()print("unrerank:")print(retriever.invoke("What is agent memory"))
llm = OpenAI(temperature=0)# 替换自己的api keycompressor = CohereRerank(cohere_api_key="xxx", top_n=10, model="rerank-english-v3.0)compression_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
compressed_docs = compression_retriever.invoke("What is agent memory")
print("============================")print("rerank:")print(compressed_docs)

from langchain.document_loaders import WebBaseLoaderfrom langchain.text_splitter import RecursiveCharacterTextSplitterfrom langchain_community.vectorstores import Chromafrom langchain_openai import OpenAIEmbeddingsfrom transformers import AutoModelForSequenceClassification, AutoTokenizerimport torch
urls = ["https://lilianweng.github.io/posts/2023-06-23-agent/",]
docs = [WebBaseLoader(url).load() for url in urls]docs_list = [item for sublist in docs for item in sublist]
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=250, chunk_overlap=0)doc_splits = text_splitter.split_documents(docs_list)

def pretty_print_docs(docs):print(f"\n{'-' * 100}\n".join([f"Document {i+1}:\n\n{d.page_content}\nMetadata: {d.metadata}"for i, d in enumerate(docs)]))

# Add to vectorDBvectorstore = Chroma.from_documents(documents=doc_splits,collection_name="rag-chroma",embedding=OpenAIEmbeddings(),)retriever = vectorstore.as_retriever()
# 加载分词器tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-large")# 加载序列化分类的模型model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-large")# 将模型设置为评估模式用于推理而非训练model.eval()device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 模型绑定cpu或者gpumodel.to(device)

def rerank(query, documents):inputs = tokenizer(list(zip([query] * len(documents),[doc.page_content for doc in documents])), padding=True, truncation=True, return_tensors='pt', max_length=512)scores = model(**inputs, return_dict=True).logits.view(-1, ).float()ranked_docs = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)return [doc for doc in ranked_docs]

def handle_query(query):results = retriever.get_relevant_documents(query)reranked_results = rerank(query, results)return reranked_results

query = "what is agent memory?"print("unrerank:")pretty_print_docs(retriever.get_relevant_documents(query))print("=" * 50)reranked_result = handle_query(query)print("rerank:")pretty_print_docs(reranked_result)

reordering = LongContextReorder()reordered_docs = reordering.transform_documents(docs)