ModelStorage (#41)

* Implemented ModelStorage class

Author: NotBioWaste905

.github/workflows/tests.yml

@@ -30,4 +30,4 @@ jobs:
 
       - name: run tests
         run: |
-          python -m poetry run poe test
+          python -m poetry run poe test || exit 1

.gitignore

@@ -10,4 +10,5 @@ __pycache__
 .pytest_cache
 .mypy_cache
 test.ipynb
-*.pyc
+*.pyc
+docs/build

dialogue2graph/cli/commands/generate_data.py

@@ -1,44 +1,21 @@
 import json
-import os
 from pathlib import Path
-from langchain_openai import ChatOpenAI
-from dialogue2graph.datasets.complex_dialogues.generation import LoopedGraphGenerator
+from dialogue2graph.pipelines.topic_generation.pipeline import TopicGenerationPipeline
+from dialogue2graph.pipelines.model_storage import ModelStorage
+
+ms = ModelStorage()
 
 
 def generate_data(topic: str, config: dict, output_path: str):
     """Generate dialogue data for a given topic"""
 
-    if config == {}:
-        gen_model = ChatOpenAI(
-            model="gpt-4o",
-            api_key=os.getenv("OPENAI_API_KEY"),
-            base_url=os.getenv("OPENAI_BASE_URL"),
-        )
-
-        val_model = ChatOpenAI(
-            model="gpt-3.5-turbo",
-            api_key=os.getenv("OPENAI_API_KEY"),
-            base_url=os.getenv("OPENAI_BASE_URL"),
-            temperature=0,
-        )
-    else:
-        gen_model = ChatOpenAI(
-            model=config["models"].get("generation-model", {}).get("name", "gpt-4o"),
-            temperature=config["models"].get("generation-model", {}).get("temperature", 0.7),
-            api_key=os.getenv("OPENAI_API_KEY"),
-            base_url=os.getenv("OPENAI_BASE_URL"),
-        )
-        val_model = ChatOpenAI(
-            model=config["models"].get("validation-model", {}).get("name", "gpt-3.5-turbo"),
-            temperature=config["models"].get("validation-model", {}).get("temperature", 0.7),
-            api_key=os.getenv("OPENAI_API_KEY"),
-            base_url=os.getenv("OPENAI_BASE_URL"),
-        )
-
-    pipeline = LoopedGraphGenerator(
-        generation_model=gen_model,
-        validation_model=val_model,
-    )
+    if config != {}:
+        ms.load(config)
+
+    pipeline = TopicGenerationPipeline(ms)
+
+    result = pipeline.invoke(topic)
+    print("Result:", result.graph_dict)
 
     result = pipeline.invoke(topic=topic)
 

dialogue2graph/cli/commands/generate_graph_algo.py

@@ -1,35 +1,19 @@
 import json
 from pathlib import Path
 from dialogue2graph.pipelines.d2g_algo.pipeline import Pipeline
-from dialogue2graph.pipelines.models import ModelsAPI
+from dialogue2graph.pipelines.model_storage import ModelStorage
 
-models = ModelsAPI()
+ms = ModelStorage()
 
 
-def generate_algo(dialogues: str, config: dict, output_path: str):
+def generate_algo(dialogues: str, config: Path, output_path: str):
     """Generates graph from dialogues via d2g_algo pipeline using parameters from config
     and saves graph dictionary to output_path"""
 
-    if config == {}:
-        filler_name = "chatgpt-4o-latest"
-        formatter_name = "gpt-4o-mini"
-        filler_temp = 0
-        formatter_temp = 0
-        sim_name = "BAAI/bge-m3"
-        device = "cpu"
-    else:
-        filler_name = config["models"].get("filler-model", {}).get("name", "chatgpt-4o-latest")
-        formatter_name = config["models"].get("formatter-model", {}).get("name", "gpt-4o-mini")
-        filler_temp = config["models"].get("filler-model", {}).get("temperature", 0)
-        formatter_temp = config["models"].get("formatter-model", {}).get("temperature", 0)
-        sim_name = config["models"].get("sim-model", {}).get("name", "BAAI/bge-m3")
-        device = config["models"].get("sim-model", {}).get("device", "cpu")
-
-    filling_llm = models("llm", name=filler_name, temp=filler_temp)
-    formatting_llm = models("llm", name=formatter_name, temp=formatter_temp)
-    sim_model = models("similarity", name=sim_name, device=device)
-
-    pipeline = Pipeline(filling_llm=filling_llm, formatting_llm=formatting_llm, sim_model=sim_model)
+    if config != {}:
+        ms.load(config)
+
+    pipeline = Pipeline(ms)
 
     result = pipeline.invoke(dialogues)
     print("Result:", result.graph_dict)

dialogue2graph/cli/commands/generate_graph_extender.py

@@ -1,40 +1,19 @@
 import json
 from pathlib import Path
 from dialogue2graph.pipelines.d2g_extender.pipeline import Pipeline
-from dialogue2graph.pipelines.models import ModelsAPI
+from dialogue2graph.pipelines.model_storage import ModelStorage
 
-models = ModelsAPI()
+ms = ModelStorage()
 
 
-def generate_extender(dialogues: str, config: dict, output_path: str):
+def generate_extender(dialogues: str, config: Path, output_path: str):
     """Generates graph from dialogues via d2g_llm pipeline using parameters from config
     and saves graph dictionary to output_path"""
 
-    if config == {}:
-        extender_name = "chatgpt-4o-latest"
-        extender_temp = 0
-        filler_name = "chatgpt-4o-latest"
-        filler_temp = 0
-        formatter_name = "gpt-4o-mini"
-        formatter_temp = 0
-        sim_name = "BAAI/bge-m3"
-        device = "cpu"
-    else:
-        extender_name = config["models"].get("extender-model", {}).get("name", "chatgpt-4o-latest")
-        extender_temp = config["models"].get("extender-model", {}).get("temperature", 0)
-        filler_name = config["models"].get("filler-model", {}).get("name", "chatgpt-4o-latest")
-        filler_temp = config["models"].get("filler-model", {}).get("temperature", 0)
-        formatter_name = config["models"].get("formatter-model", {}).get("name", "gpt-4o-mini")
-        formatter_temp = config["models"].get("formatter-model", {}).get("temperature", 0)
-        sim_name = config["models"].get("sim-model", {}).get("name", "BAAI/bge-m3")
-        device = config["models"].get("sim-model", {}).get("device", "cpu")
-
-    extending_llm = models("llm", name=extender_name, temp=extender_temp)
-    filling_llm = models("llm", name=filler_name, temp=filler_temp)
-    formatting_llm = models("llm", name=formatter_name, temp=formatter_temp)
-    sim_model = models("similarity", name=sim_name, device=device)
-
-    pipeline = Pipeline(extending_llm=extending_llm, filling_llm=filling_llm, formatting_llm=formatting_llm, sim_model=sim_model)
+    if config != {}:
+        ms.load(config)
+
+    pipeline = Pipeline(ms)
 
     result = pipeline.invoke(dialogues)
     print("Result:", result.graph_dict)

dialogue2graph/cli/commands/generate_graph_llm.py

@@ -1,40 +1,18 @@
 import json
 from pathlib import Path
 from dialogue2graph.pipelines.d2g_llm.pipeline import Pipeline
-from dialogue2graph.pipelines.models import ModelsAPI
+from dialogue2graph.pipelines.model_storage import ModelStorage
 
-models = ModelsAPI()
+ms = ModelStorage()
 
 
-def generate_llm(dialogues: str, config: dict, output_path: str):
+def generate_llm(dialogues: str, config: Path, output_path: str):
     """Generates graph from dialogues via d2g_llm pipeline using parameters from config
     and saves graph dictionary to output_path"""
 
-    if config == {}:
-        grouper_name = "chatgpt-4o-latest"
-        grouper_temp = 0
-        filler_name = "chatgpt-4o-latest"
-        filler_temp = 0
-        formatter_name = "gpt-4o-mini"
-        formatter_temp = 0
-        sim_name = "BAAI/bge-m3"
-        device = "cpu"
-    else:
-        grouper_name = config["models"].get("grouper-model", {}).get("name", "chatgpt-4o-latest")
-        grouper_temp = config["models"].get("grouper-model", {}).get("temperature", 0)
-        filler_name = config["models"].get("filler-model", {}).get("name", "chatgpt-4o-latest")
-        filler_temp = config["models"].get("filler-model", {}).get("temperature", 0)
-        formatter_name = config["models"].get("formatter-model", {}).get("name", "gpt-4o-mini")
-        formatter_temp = config["models"].get("formatter-model", {}).get("temperature", 0)
-        sim_name = config["models"].get("sim-model", {}).get("name", "BAAI/bge-m3")
-        device = config["models"].get("sim-model", {}).get("device", "cpu")
-
-    grouping_llm = models("llm", name=grouper_name, temp=grouper_temp)
-    filling_llm = models("llm", name=filler_name, temp=filler_temp)
-    formatting_llm = models("llm", name=formatter_name, temp=formatter_temp)
-    sim_model = models("similarity", name=sim_name, device=device)
-
-    pipeline = Pipeline(grouping_llm=grouping_llm, filling_llm=filling_llm, formatting_llm=formatting_llm, sim_model=sim_model)
+    if config != {}:
+        ms.load(config)
+    pipeline = Pipeline(ms)
 
     result = pipeline.invoke(dialogues)
     print("Result:", result.graph_dict)

dialogue2graph/cli/main.py

@@ -1,5 +1,4 @@
 import click
-import yaml
 from dotenv import load_dotenv
 from .commands.generate_data import generate_data
 from .commands.generate_graph_algo import generate_algo
@@ -21,14 +20,7 @@ def cli():
 def gen_data(env: str, cfg: str, topic: str, output: str):
     """Generate dialogue data for a given topic"""
     load_dotenv(env)
-    with open(cfg) as stream:
-        config: dict = {}
-        try:
-            config = yaml.safe_load(stream)
-        except yaml.YAMLError as exc:
-            print(exc)
-
-    generate_data(topic, config, output)
+    generate_data(topic, cfg, output)
 
 
 @cli.command()
@@ -39,13 +31,7 @@ def gen_data(env: str, cfg: str, topic: str, output: str):
 def gen_graph_algo(env: str, cfg: str, dialogues: str, output: str):
     """Generate graph from dialogues data via d2g_algo pipeline"""
     load_dotenv(env)
-    with open(cfg) as stream:
-        config: dict = {}
-        try:
-            config = yaml.safe_load(stream)
-        except yaml.YAMLError as exc:
-            print(exc)
-    generate_algo(dialogues, config, output)
+    generate_algo(dialogues, cfg, output)
 
 
 @cli.command()
@@ -56,13 +42,7 @@ def gen_graph_algo(env: str, cfg: str, dialogues: str, output: str):
 def gen_graph_llm(env: str, cfg: str, dialogues: str, output: str):
     """Generate graph from dialogues data via d2g_llm pipeline"""
     load_dotenv(env)
-    with open(cfg) as stream:
-        config: dict = {}
-        try:
-            config = yaml.safe_load(stream)
-        except yaml.YAMLError as exc:
-            print(exc)
-    generate_llm(dialogues, config, output)
+    generate_llm(dialogues, cfg, output)
 
 
 @cli.command()
@@ -73,13 +53,7 @@ def gen_graph_llm(env: str, cfg: str, dialogues: str, output: str):
 def gen_graph_extender(env: str, cfg: str, dialogues: str, output: str):
     """Generate graph from dialogues data via d2g_llm pipeline"""
     load_dotenv(env)
-    with open(cfg) as stream:
-        config: dict = {}
-        try:
-            config = yaml.safe_load(stream)
-        except yaml.YAMLError as exc:
-            print(exc)
-    generate_extender(dialogues, config, output)
+    generate_extender(dialogues, cfg, output)
 
 
 if __name__ == "__main__":

dialogue2graph/datasets/__init__.py

@@ -0,0 +1 @@
+from dialogue2graph.datasets.core import Dataset

dialogue2graph/datasets/complex_dialogues/generation.py

@@ -15,6 +15,7 @@
 from dialogue2graph.pipelines.core.graph import BaseGraph, Graph
 from dialogue2graph.pipelines.core.algorithms import TopicGraphGenerator
 from dialogue2graph.pipelines.core.schemas import GraphGenerationResult, DialogueGraph
+from dialogue2graph.pipelines.model_storage import ModelStorage
 from dialogue2graph.utils.prompt_caching import setup_cache, add_uuid_to_prompt
 
 from .prompts import cycle_graph_generation_prompt_informal, cycle_graph_repair_prompt, graph_example
@@ -276,19 +277,30 @@ def __call__(self, topic: str) -> PipelineResult:
 
 
 class LoopedGraphGenerator(TopicGraphGenerator):
-    generation_model: BaseChatModel
-    validation_model: BaseChatModel
+    """Graph generator for topic-based dialogue generation with model storage support"""
+
+    model_storage: ModelStorage = Field(description="Model storage")
+    generation_llm: str = Field(description="LLM for graph generation")
+    validation_llm: str = Field(description="LLM for validation")
+    theme_validation_llm: str = Field(description="LLM for theme validation")
     pipeline: GenerationPipeline
 
-    def __init__(self, generation_model: BaseChatModel, validation_model: BaseChatModel, theme_validation_model: BaseChatModel):
+    def __init__(
+        self,
+        model_storage: ModelStorage,
+        generation_llm: str,
+        validation_llm: str,
+        theme_validation_llm: str,
+    ):
         super().__init__(
-            generation_model=generation_model,
-            validation_model=validation_model,
-            theme_validation_model=theme_validation_model,
+            model_storage=model_storage,
+            generation_llm=generation_llm,
+            validation_llm=validation_llm,
+            theme_validation_llm=theme_validation_llm,
             pipeline=GenerationPipeline(
-                generation_model=generation_model,
-                validation_model=validation_model,
-                theme_validation_model=theme_validation_model,
+                generation_model=model_storage.storage[generation_llm].model,
+                validation_model=model_storage.storage[validation_llm].model,
+                theme_validation_model=model_storage.storage[theme_validation_llm].model,
                 generation_prompt=cycle_graph_generation_prompt_informal,
                 repair_prompt=cycle_graph_repair_prompt,
             ),

dialogue2graph/datasets/core/__init__.py

@@ -0,0 +1,3 @@
+from .dataset import Dataset
+
+__all__ = ["Dataset"]