Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 6,471 Bytes
6e20157 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
import os
import random
from functools import cache
from operator import itemgetter
import langsmith
from langchain.memory import ConversationBufferWindowMemory
from langchain.retrievers import EnsembleRetriever
from langchain_community.document_transformers import LongContextReorder
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableLambda
from langchain_openai.chat_models import ChatOpenAI
from .prompt_template import generate_prompt_template
from .retrievers_setup import (DenseRetrieverClient, SparseRetrieverClient,
compression_retriever_setup)
# Helpers
def reorder_documents(docs: list[Document]) -> list[Document]:
"""Long-Context Reorder: No matter the architecture of the model, there is
a performance degradation when we include 10+ retrieved documents.
Args:
docs (list): List of Langchain documents
Returns:
list: Reordered list of Langchain documents
"""
reorder = LongContextReorder()
return reorder.transform_documents(docs)
def randomize_documents(documents: list[Document]) -> list[Document]:
"""Randomize the documents to vary the recommendations."""
random.shuffle(documents)
return documents
def format_practitioners_docs(docs: list[Document]) -> str:
"""Format the practitioners_db Documents to markdown.
Args:
docs (list[Documents]): List of Langchain documents
Returns:
docs (str):
"""
return f"\n{'-' * 3}\n".join(
[f"- Practitioner #{i+1}:\n\n\t" +
d.page_content for i, d in enumerate(docs)]
)
def format_tall_tree_docs(docs: list[Document]) -> str:
"""Format the tall_tree_db Documents to markdown.
Args:
docs (list[Documents]): List of Langchain documents
Returns:
docs (str):
"""
return f"\n{'-' * 3}\n".join(
[f"- No. {i+1}:\n\n\t" +
d.page_content for i, d in enumerate(docs)]
)
def create_langsmith_client():
"""Create a Langsmith client."""
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "talltree-ai-assistant"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
langsmith_api_key = os.getenv("LANGCHAIN_API_KEY")
if not langsmith_api_key:
raise EnvironmentError(
"Missing environment variable: LANGCHAIN_API_KEY")
return langsmith.Client()
# Set up Runnable and Memory
@cache
def get_rag_chain(model_name: str = "gpt-4", temperature: float = 0.2) -> tuple[ChatOpenAI, ConversationBufferWindowMemory]:
"""Set up runnable and chat memory
Args:
model_name (str, optional): LLM model. Defaults to "gpt-4" 30012024.
temperature (float, optional): Model temperature. Defaults to 0.2.
Returns:
Runnable, Memory: Chain and Memory
"""
# Set up Langsmith to trace the chain
langsmith_tracing = create_langsmith_client()
# LLM and prompt template
llm = ChatOpenAI(model_name=model_name,
temperature=temperature)
prompt = generate_prompt_template()
# Set retrievers pointing to the practitioners's dataset
embeddings_model = "text-embedding-ada-002"
dense_retriever_client = DenseRetrieverClient(embeddings_model=embeddings_model,
collection_name="practitioners_db")
# Qdrant db as a retriever
practitioners_db_dense_retriever = dense_retriever_client.get_dense_retriever(search_type="similarity",
k=10)
# Testing the sparse vector retriever using Qdrant
collection_name = "practitioners_db_sparse_collection"
vector_name = "sparse_vector"
sparse_retriever_client = SparseRetrieverClient(
collection_name=collection_name,
vector_name=vector_name,
splade_model_id="naver/splade-cocondenser-ensembledistil",
k=15)
practitioners_db_sparse_retriever = sparse_retriever_client.get_sparse_retriever()
# TODO Test the ensemble retriever for hyprid search (dense retriever seems to work better)
# Using only the filtered sparse retriever
practitioners_ensemble_retriever = EnsembleRetriever(
retrievers=[practitioners_db_dense_retriever,
practitioners_db_sparse_retriever], weights=[0.1, 0.9]
)
# Compression retriever for practitioners db
# TODO Test the filtered ensemble retriever *** Using only the sparse retriever ***
practitioners_db_compression_retriever = compression_retriever_setup(
practitioners_db_sparse_retriever,
embeddings_model="text-embedding-ada-002",
similarity_threshold=0.74
)
# Set retrievers pointing to the tall_tree_db
dense_retriever_client = DenseRetrieverClient(embeddings_model=embeddings_model,
collection_name="tall_tree_db")
tall_tree_db_dense_retriever = dense_retriever_client.get_dense_retriever(search_type="similarity",
k=5)
# Compression retriever for tall_tree_db
tall_tree_db_compression_retriever = compression_retriever_setup(
tall_tree_db_dense_retriever,
embeddings_model="text-embedding-ada-002",
similarity_threshold=0.5
)
# Set conversation history window memory. It only uses the last k=4 interactions.
memory = ConversationBufferWindowMemory(memory_key="history",
return_messages=True,
k=5)
# Set up runnable using LCEL
setup_and_retrieval = {"practitioners_db": itemgetter("message")
| practitioners_db_compression_retriever
| randomize_documents
| format_practitioners_docs,
"tall_tree_db": itemgetter("message") | tall_tree_db_compression_retriever | format_tall_tree_docs,
"history": RunnableLambda(memory.load_memory_variables) | itemgetter("history"),
"message": itemgetter("message")
}
chain = (
setup_and_retrieval
| prompt
| llm
| StrOutputParser()
)
return chain, memory
|