feat: add configurable

config.yml

@@ -0,0 +1,14 @@
+models:
+  casual_conversation:
+    model: openchat/openchat-7b
+    temperature: 0.7
+  multiquery_retrieval:
+    model: openchat/openchat-7b
+    temperature: 0.3
+  classificator_msg:
+    model: openchat/openchat-7b
+    temperature: 0
+  rag:
+    model: cohere/command-r
+    temperature: 0
+  

lib/config.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+
+import yaml
+
+
+class Config:
+    def __init__(self, config_path: Path):
+        self.config_path = config_path
+        self._config = self._load_config()
+
+    def _load_config(self) -> dict:
+        with open(self.config_path, "r") as file:
+            return yaml.safe_load(file)
+
+    def get(self, *keys, default=None):
+        config = self._config
+        for key in keys:
+            config = config.get(key, default)
+            if config is default:
+                break
+        return config
+
+    def __getitem__(self, item: str):
+        return self._config.get(item)
+
+    def __repr__(self) -> str:
+        return f"Config({self._config})"

lib/graph.py

@@ -15,6 +15,25 @@ from lib.runnables import (
     message_classificator,
 )
 from langgraph.graph import END, StateGraph
+from transformers import AutoModel, AutoTokenizer
+import torch
+import torch.nn.functional as F
+
+model_name = "intfloat/multilingual-e5-large"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+device = torch.device("cuda")
+model.to(device)
+SIMILARITY_TRESHHOLD = 0.8
+
+
+def get_embeddings(texts):
+    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
+    inputs.to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        embeddings = torch.mean(outputs.last_hidden_state, dim=1)
+    return embeddings
 
 
 class GraphState(TypedDict):
@@ -32,7 +51,6 @@ class GraphState(TypedDict):
     generation: str
     documents: List[Document]
     filtered_documets: List[Document]
-    is_fuse: bool
     count_regenerations: int
 
 
@@ -162,6 +180,40 @@ def generate(state):
     return {"documents": documents, "question": question, "generation": generation}
 
 
+def grade_documents_by_embed(state):
+    """
+    Determines whether the retrieved documents are relevant to the question.
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): Updates documents key with only filtered relevant documents
+    """
+    question = state["question"]
+    documents = state["documents"]
+
+    # Score each doc
+    filtered_docs = []
+
+    query_embedding = get_embeddings([question])
+    document_embeddings = get_embeddings([doc.page_content for doc in documents])
+
+    # Calculate cosine similarity
+    similarity_scores = F.cosine_similarity(query_embedding, document_embeddings)
+    for doc, score in zip(documents, similarity_scores):
+        if score >= SIMILARITY_TRESHHOLD:
+            filtered_docs.append(doc)
+    sorted_documents = [
+        doc[0]
+        for doc in sorted(
+            zip(documents, similarity_scores), key=lambda x: x[1], reverse=True
+        )
+    ]
+    cut_off_documents = sorted_documents[:5]
+    return {"documents": cut_off_documents, "question": question}
+
+
 def grade_documents(state):
     """
     Determines whether the retrieved documents are relevant to the question.
@@ -199,7 +251,6 @@ def grade_documents(state):
                 filtered_docs.append(documents[ind_d + j])
             else:
                 print("---GRADE: DOCUMENT NOT RELEVANT---")
-    is_fuse = len(filtered_docs) / len(documents) <= 0.5
 
     return {"documents": filtered_docs, "question": question}
 
@@ -264,7 +315,7 @@ def build_workflow():
     # Define the nodes
     workflow.add_node("start_point", start_point)
     workflow.add_node("retrieve", retrieve)  # retrieve
-    workflow.add_node("grade_documents", grade_documents)  # grade documents
+    workflow.add_node("grade_documents", grade_documents_by_embed)  # grade documents
     workflow.add_node("generate", generate)  # generate
     workflow.add_node("casual_chat", casual_chat)  # simple chat
     workflow.add_node("add_sources", add_sources)

lib/prompts.py

@@ -29,11 +29,10 @@ rag_assistant_prompt = PromptTemplate(
     template="""
         SYSTEM: You are an assistant for question-answering tasks. 
         Use the following pieces of retrieved context to answer the question.
-        Use previous messages then current message higly likely  
+        Use a 'Previous messages' as a part of context.
         If you don't find the answer in the context, transform the question ans ask the user to specify his qusetion. 
 
         Keep the answer concise.  
-        Print a most possible topic of conversation.
         Always reply in Russian, all text must be in Russian!
 
         Context: {context}

lib/runnables.py

@@ -1,6 +1,7 @@
 import contextlib
 
 from lib.model_builder import ModelBuilderV2
+from lib.config import Config
 from lib.prompts import (
     casual_prompt,
     grader_3_doc_prompt,
@@ -16,6 +17,7 @@ from langchain_core.chat_history import (
 from langchain_core.runnables import ConfigurableFieldSpec
 
 store = {}
+config = Config("config.yml")
 
 
 def get_session_history(user_id: str, conversation_id: str) -> BaseChatMessageHistory:
@@ -25,9 +27,9 @@ def get_session_history(user_id: str, conversation_id: str) -> BaseChatMessageHi
 
 
 class ModelConfig:
-    def __init__(self, model_name, temperature=0.7):
-        self.model_name = model_name
-        self.temperature = temperature
+    def __init__(self, config_key):
+        self.model_name = config.get("models", config_key, "model")
+        self.temperature = config.get("models", config_key, "temperature")
 
 
 class ConfigField:
@@ -81,10 +83,10 @@ def create_model_builder(config):
     #     llm.release()  # Assuming ModelBuilderV2 has a release method to clear resources
 
 
-casual_config = ModelConfig("openchat/openchat-7b", 0.7)
-retrieval_config = ModelConfig("cohere/command-r")
-rag_config = ModelConfig("mistralai/mixtral-8x22b-instruct")
-classificator_msg_config = ModelConfig("openchat/openchat-7b")
+casual_config = ModelConfig("casual_conversation")
+multiquery_config = ModelConfig("multiquery_retrieval")
+rag_config = ModelConfig("rag")
+classificator_msg_config = ModelConfig("classificator_msg")
 
 history_config = [USER_ID_FIELD, CONVERSATION_ID_FIELD]
 
@@ -96,7 +98,7 @@ with create_model_builder(casual_config) as llm:
         | StrOutputParser()
     )
 
-with create_model_builder(retrieval_config) as llm:
+with create_model_builder(multiquery_config) as llm:
     retrieval_grader_3 = grader_3_doc_prompt | llm | JsonOutputParser()
 
 with create_model_builder(rag_config) as llm: