refactoring to easly use

README.md

@@ -12,8 +12,8 @@ In recent years, various LLMs, embedding models, and LLM flows utilizing them ha
 
 This repository aims to treat LLMs and Embeddings as a hyperparameter, with the goal of automatically searching for the optimal hyperparameter of the LLM flow.
 
-**Below image is the concept image of this repository (image is took from [Flowise](https://github.com/FlowiseAI/Flowise) and slightly modified). Component of LangChain like a LLM or Embedding can be treated as hyperparameter. You will find component from various candidate that can optimize score.**
 ![concept_image](documents/image/concept.png)
+**This image is the concept image of this repository (image is took from [Flowise](https://github.com/FlowiseAI/Flowise) and slightly modified). Component of LangChain like a LLM or Embedding can be treated as hyperparameter. You will find component from various candidate that can optimize score.**
 
 This repository is strongly inspired by [lightning-hydra-template](https://github.com/ashleve/lightning-hydra-template)🎉
 

llmflowoptimizer/component/evaluation/sample.py

@@ -12,10 +12,8 @@
 )
 from scipy import stats
 
-from llmflowoptimizer.component.base.base import BaseChainModel, BaseEvaluationModel
 
-
-class Evaluation(BaseEvaluationModel):
+class Evaluation:
     """Define the evaluation system. llmflowoptimizer optimizes the hyperparameters of the model
     based on the output of this evaluation system.
 
@@ -58,14 +56,14 @@ def __init__(
 
     def evaluate(
         self,
-        model: BaseChainModel,
+        model: Any,  # this model should be defined in llmflowoptimizer/component/model/sample_qa.py
     ):
         """Return of this method is used for hyperparameter optimization."""
         project_name = self.project_name + datetime.datetime.now().strftime("_%Y-%m-%d_%H-%M-%S")
         result = run_on_dataset(
             client=self.client,
             dataset_name=self.dataset_name,
-            llm_or_chain_factory=model,
+            llm_or_chain_factory=model.get_chain(),
             evaluation=self.evaluation_config,
             project_name=project_name,
             input_mapper=lambda x: x,

llmflowoptimizer/component/model/sample_qa.py

@@ -6,10 +6,8 @@
 from langchain.schema.language_model import BaseLanguageModel
 from langchain.text_splitter import TextSplitter
 
-from llmflowoptimizer.component.base.base import BaseChainModel
 
-
-class SampleQA(BaseChainModel):
+class SampleQA:
     """Define the flow of the model to be adjusted."""
 
     def __init__(

llmflowoptimizer/run.py

@@ -1,12 +1,10 @@
 import logging
-from pprint import pprint
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Optional
 
 import hydra
 import rootutils
 from omegaconf import DictConfig
 
-from llmflowoptimizer.component.base.base import BaseChainModel, BaseEvaluationModel
 from llmflowoptimizer.utils.utils import print_config_tree
 
 log = logging.getLogger(__name__)
@@ -19,14 +17,13 @@
 def main(cfg: DictConfig) -> Optional[float]:
     """This is the main entry point of the tune script."""
     log.info(f"Instantiating model <{cfg.model._target_}>")
-    model: BaseChainModel = hydra.utils.instantiate(cfg.model)
-    model = model.get_chain()
+    model = hydra.utils.instantiate(cfg.model)
 
     if cfg.extras.print_config:
         print_config_tree(cfg)
 
     if cfg.extras.evaluation:
-        evaluator: BaseEvaluationModel = hydra.utils.instantiate(cfg.evaluation)
+        evaluator = hydra.utils.instantiate(cfg.evaluation)
         log.info(f"Evaluating <{cfg.model._target_}>")
         metric_value = evaluator.evaluate(model)
         log.info(f"Score is  {metric_value}")