File size: 653 Bytes
7e4014b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import ChatOpenAI
from ragatouille import RAGPretrainedModel
def create_retriever(texts):
colbert = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv1.9")
colbert.index(
collection=[chunk.page_content for chunk in texts],
split_documents=False,
document_metadatas=[chunk.metadata for chunk in texts],
index_name="vector_store",
)
retriever = colbert.as_langchain_retriever(k=5)
retriever = MultiQueryRetriever.from_llm(
retriever=retriever, llm=ChatOpenAI(temperature=0)
)
return retriever
|