Feat/model validation (#48)

* Initializing models outside of the ModelStorage

* ModelStorage.add() now demand type instead of string in model_type field

* update userguides

---------

Co-authored-by: Yuriy Peshkichev <[email protected]>
Co-authored-by: NotBioWaste905 <[email protected]>

Author: 3 people

dialogue2graph/datasets/complex_dialogues/generation.py

@@ -5,7 +5,6 @@
 The module provides graph generator capable of creating complex validated graphs.
 """
 
-import logging
 import os
 from enum import Enum
 from typing import Optional, Dict, Any, Union
@@ -37,8 +36,9 @@
 )
 
 # Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from dialogue2graph.utils.logger import Logger
+
+logger = Logger(__file__)
 
 
 class ErrorType(str, Enum):

dialogue2graph/metrics/llm_metrics/metrics.py

@@ -5,7 +5,6 @@
 The module contains functions that checks Graphs and Dialogues for various metrics using LLM calls.
 """
 
-import logging
 import json
 from typing import List, TypedDict, Union
 from pydantic import BaseModel, Field
@@ -22,8 +21,9 @@
 from langchain.chat_models import ChatOpenAI
 from langchain.schema import HumanMessage
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+from dialogue2graph.utils.logger import Logger
+
+logger = Logger(__file__)
 
 
 class InvalidTransition(TypedDict):

dialogue2graph/metrics/llm_validators/validators.py

@@ -13,7 +13,9 @@
 from dialogue2graph.pipelines.model_storage import ModelStorage
 from dialogue2graph.metrics.similarity import compare_strings
 
-from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_openai import ChatOpenAI
+from langchain_core.language_models import BaseChatModel
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain.output_parsers import PydanticOutputParser
 
@@ -125,7 +127,10 @@ def is_greeting_repeated_emb_llm(
         starts = START_TURNS
 
     if model_storage.storage.get(embedder_name):
-        if not model_storage.storage.get(embedder_name).model_type == "emb":
+        if (
+            not model_storage.storage.get(embedder_name).model_type
+            == HuggingFaceEmbeddings
+        ):
             raise TypeError(f"The {embedder_name} model is not an embedder")
         embedder_model = model_storage.storage[embedder_name].model
     else:
@@ -134,7 +139,7 @@ def is_greeting_repeated_emb_llm(
         )
 
     if model_storage.storage.get(llm_name):
-        if not model_storage.storage.get(llm_name).model_type == "llm":
+        if not model_storage.storage.get(llm_name).model_type == ChatOpenAI:
             raise TypeError(f"The {llm_name} model is not an LLM")
         llm_model = model_storage.storage[llm_name].model
     else:
@@ -183,7 +188,10 @@ def is_dialog_closed_too_early_emb_llm(
         ends = END_TURNS
 
     if model_storage.storage.get(embedder_name):
-        if not model_storage.storage.get(embedder_name).model_type == "emb":
+        if (
+            not model_storage.storage.get(embedder_name).model_type
+            == HuggingFaceEmbeddings
+        ):
             raise TypeError(f"The {embedder_name} model is not an embedder")
         embedder_model = model_storage.storage[embedder_name].model
     else:
@@ -192,7 +200,7 @@ def is_dialog_closed_too_early_emb_llm(
         )
 
     if model_storage.storage.get(llm_name):
-        if not model_storage.storage.get(llm_name).model_type == "llm":
+        if not model_storage.storage.get(llm_name).model_type == ChatOpenAI:
             raise TypeError(f"The {llm_name} model is not an LLM")
         llm_model = model_storage.storage[llm_name].model
     else:

dialogue2graph/metrics/no_llm_metrics/metrics.py

@@ -13,6 +13,9 @@
 
 from dialogue2graph.pipelines.core.graph import BaseGraph
 from dialogue2graph.pipelines.core.dialogue import Dialogue
+from dialogue2graph.utils.logger import Logger
+
+logger = Logger(__file__)
 
 
 logging.basicConfig(level=logging.INFO)

dialogue2graph/pipelines/core/dialogue_sampling.py

@@ -6,7 +6,6 @@
 """
 
 import itertools
-import logging
 from typing import Literal
 import pandas as pd
 from dialogue2graph.pipelines.core.graph import BaseGraph
@@ -19,8 +18,9 @@
 from dialogue2graph.pipelines.helpers.find_cycle_ends import find_cycle_ends
 from langchain_core.language_models.chat_models import BaseChatModel
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from dialogue2graph.utils.logger import Logger
+
+logger = Logger(__file__)
 
 
 class _DialogPathsCounter:
@@ -189,23 +189,21 @@ def remove_duplicated_paths(node_paths: list[list[int]]) -> list[list[int]]:
     return res
 
 
-def get_dialogue_doublets(seq: list[list[dict]]) -> set[tuple[str]]:
-    """Find all dialogue doublets with (edge, target) utterances
-
-    Args:
-        seq: sequence of dialogs
-
-    Returns:
-        Set of (user_utterance, assistant_utterance)
-    """
-    doublets = set()
-    for dialogue in seq:
-        user_texts = [d["text"] for d in dialogue if d["participant"] == "user"]
-        assist_texts = [d["text"] for d in dialogue if d["participant"] == "assistant"]
-        if len(assist_texts) > len(user_texts):
-            user_texts += [""]
-        doublets.update(zip(user_texts, assist_texts))
-    return doublets
+# def get_dialogue_doublets(seq: list[list[dict]]) -> set[tuple[str]]:
+#     """Find all dialogue doublets with (edge, target) utterances
+#     Args:
+#       seq: sequence of dialogs
+#     Returns:
+#       Set of (user_utterance, assistant_utterance)
+#     """
+#     doublets = set()
+#     for dialogue in seq:
+#         user_texts = [d["text"] for d in dialogue if d["participant"] == "user"]
+#         assist_texts = [d["text"] for d in dialogue if d["participant"] == "assistant"]
+#         if len(assist_texts) > len(user_texts):
+#             user_texts += [""]
+#         doublets.update(zip(user_texts, assist_texts))
+#     return doublets
 
 
 def get_dialogue_triplets(seq: list[list[dict]]) -> set[tuple[str]]:
@@ -239,9 +237,7 @@ def remove_duplicated_dialogues(seq: list[list[dict]]) -> list[list[dict]]:
         return []
     uniq_seq = [non_empty_seq[0]]
     for s in non_empty_seq[1:]:
-        if not get_dialogue_doublets([s]).issubset(
-            get_dialogue_doublets(uniq_seq)
-        ) or not get_dialogue_triplets([s]).issubset(get_dialogue_triplets(uniq_seq)):
+        if not get_dialogue_triplets([s]).issubset(get_dialogue_triplets(uniq_seq)):
             uniq_seq.append(s)
     return uniq_seq
 

dialogue2graph/pipelines/core/graph.py

@@ -10,9 +10,10 @@
 from typing import Optional, Any
 import matplotlib.pyplot as plt
 import abc
-import logging
 
-logger = logging.getLogger(__name__)
+from dialogue2graph.utils.logger import Logger
+
+logger = Logger(__file__)
 
 
 class BaseGraph(BaseModel, abc.ABC):
@@ -140,13 +141,13 @@ def load_graph(self):
         """
         self.graph = nx.DiGraph()
         nodes = sorted([v["id"] for v in self.graph_dict["nodes"]])
-        logging.debug(f"Nodes: {nodes}")
+        logger.debug(f"Nodes: {nodes}")
 
         self.node_mapping = {}
         renumber_flg = nodes != list(range(1, len(nodes) + 1))
         if renumber_flg:
             self.node_mapping = {node_id: idx + 1 for idx, node_id in enumerate(nodes)}
-        logging.debug(f"Renumber flag: {renumber_flg}")
+        logger.debug(f"Renumber flag: {renumber_flg}")
 
         for node in self.graph_dict["nodes"]:
             cur_node_id = node["id"]

dialogue2graph/pipelines/d2g_extender/pipeline.py

@@ -7,9 +7,12 @@
 
 from typing import Callable
 from dotenv import load_dotenv
+
 from dialogue2graph.pipelines.core.pipeline import BasePipeline
 from dialogue2graph.pipelines.model_storage import ModelStorage
 from dialogue2graph.pipelines.d2g_extender.three_stages_extender import LLMGraphExtender
+from langchain_openai import ChatOpenAI
+from langchain_huggingface import HuggingFaceEmbeddings
 
 load_dotenv()
 
@@ -32,6 +35,37 @@ def __init__(
         end_evals: list[Callable] = None,
         step: int = 2,
     ):
+        # if model is not in model storage put the default model there
+        model_storage.add(
+            key=extending_llm,
+            config={"model_name": "chatgpt-4o-latest", "temperature": 0},
+            model_type="llm",
+        )
+
+        model_storage.add(
+            key=filling_llm,
+            config={"mode_name": "o3-mini", "temperature": 1},
+            model_type=ChatOpenAI,
+        )
+
+        model_storage.add(
+            key=formatting_llm,
+            config={"model_name": "gpt-4o-mini", "temperature": 0},
+            model_type=ChatOpenAI,
+        )
+
+        model_storage.add(
+            key=dialog_llm,
+            config={"model_name": "o3-mini", "temperature": 1},
+            model_type=ChatOpenAI,
+        )
+
+        model_storage.add(
+            key=sim_model,
+            config={"model_name": "BAAI/bge-m3", "device": "cpu"},
+            model_type=HuggingFaceEmbeddings,
+        )
+
         super().__init__(
             name=name,
             steps=[

dialogue2graph/pipelines/d2g_extender/three_stages_extender.py

@@ -13,6 +13,7 @@
 from langchain.schema import HumanMessage
 from langchain.prompts import PromptTemplate
 
+from dialogue2graph.utils.logger import Logger
 from dialogue2graph import metrics
 from dialogue2graph.pipelines.core.dialogue_sampling import RecursiveDialogueSampler
 from dialogue2graph.pipelines.d2g_light.three_stages_light import LightGraphGenerator
@@ -44,6 +45,8 @@ class DialogueNodes(BaseModel):
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 logging.getLogger("langchain_core.vectorstores.base").setLevel(logging.ERROR)
+logger = Logger(__file__)
+
 dialogue_sampler = RecursiveDialogueSampler()
 
 

dialogue2graph/pipelines/d2g_light/pipeline.py

@@ -10,6 +10,8 @@
 from dialogue2graph.pipelines.core.pipeline import BasePipeline
 from dialogue2graph.pipelines.d2g_light.three_stages_light import LightGraphGenerator
 from dialogue2graph.pipelines.model_storage import ModelStorage
+from langchain_openai import ChatOpenAI
+from langchain_huggingface import HuggingFaceEmbeddings
 
 load_dotenv()
 
@@ -27,6 +29,22 @@ def __init__(
         step2_evals: list[Callable] = None,
         end_evals: list[Callable] = None,
     ):
+        # if model is not in model storage put the default model there
+        model_storage.add(
+            key=filling_llm,
+            config={"model_name": "chatgpt-4o-latest", "temperature": 0},
+            model_type=ChatOpenAI,
+        )
+        model_storage.add(
+            key=formatting_llm,
+            config={"model_name": "gpt-4o-mini", "temperature": 0},
+            model_type=ChatOpenAI,
+        )
+        model_storage.add(
+            key=sim_model,
+            config={"model_name": "BAAI/bge-m3", "model_kwargs": {"device": "cpu"}},
+            model_type=HuggingFaceEmbeddings,
+        )
         super().__init__(
             name=name,
             steps=[

dialogue2graph/pipelines/d2g_light/three_stages_light.py

@@ -69,7 +69,6 @@ def __init__(
         filling_llm: str = "three_stages_light_filling_llm:v1",
         formatting_llm: str = "three_stages_light_formatting_llm:v1",
         sim_model: str = "three_stages_light_sim_model:v1",
-
         step2_evals: list[Callable] | None = [],
         end_evals: list[Callable] | None = [],
     ):