Merge remote-tracking branch 'origin/dataset-generator' into DAGE-18/Persistent-Storage-JSON

Author: Intrinsical-AI

rre-dataset-generator/README.md

@@ -45,6 +45,9 @@ output_destination: "output/generated_dataset.json"
 
 ### 3. Running the Generator
 
+Before running the command below, you need to have running search engine instance (`solr`/`opensearch`/`elasticsearch`/`vespa`).
+
+
 Execute the main script via the `argparse` CLI, pointing to your configuration file:
 ```bash
 uv run dataset_generator.py --config_file config.yaml
@@ -98,4 +101,28 @@ docker compose up --build
 
 This will start 2 services:
  - `solr`, available at http://localhost:8983/solr
- - `solr-init`, loads documents from solr/data/dataset.json only if Solr doesn't index any documents.
+ - `solr-init`, loads documents from solr-init/data/dataset.json.
+
+
+##### Running OpenSearch (Single Node)
+
+To run a local OpenSearch test environment using docker-compose:
+```bash
+cd tests/integration/
+```
+
+Depending on your Docker version, you may need to use `docker compose` instead of `docker-compose`.
+If you have Docker Compose v1 installed, use:
+
+```bash
+docker-compose -f docker-compose.opensearch.yml up --build
+```
+If you have Docker Compose v2 installed, use:
+```bash
+docker compose -f docker-compose.opensearch.yml up --build
+```
+
+This will start 2 services:
+ - `opensearch`, available at http://localhost:9200/
+ - `opensearch-init`, loads documents (`bulk indexing`) from opensearch-init/data/dataset.jsonl.
+

rre-dataset-generator/config.yaml

@@ -2,13 +2,14 @@
 # dataset-generator module
 # Default: "q=#$query##"
 # Example (solr):
-query_template: "q=#$query##&df=description&rows=3&q.op=OR&wt=json"
+query_template: "q=#$query##&df=description&rows=2&q.op=OR&wt=json"
 
 # The type of search engine to use
 # Accepted values: solr, elasticsearch, opensearch, vespa
 search_engine_type: "solr"
 
 # Endpoint URL of the search engine collection or index
+# opensearch: http://localhost:9200/testcore/
 search_engine_collection_endpoint: "http://localhost:8983/solr/testcore/"
 
 # (Optional) Filter query to restrict the set of documents used to generate queries.

rre-dataset-generator/dataset_generator.py

@@ -1,13 +1,19 @@
 # configuration params
+from src.config import Config
+from src.utils import parse_args
+from src.llm.llm_config import LLMConfig
+
+# data types
+from typing import List
 from langchain_core.language_models import BaseChatModel
+from src.llm.llm_service import LLMService
+from src.model.document import Document
 from src.writers.abstract_writer import AbstractWriter
 from src.search_engine.search_engine_base import BaseSearchEngine
-from src.utils import parse_args
-from src.config import Config
-from src.llm.llm_service import LLMService
-from src.llm.llm_config import LLMConfig
+from src.model.query_response import LLMQueryResponse
+from src.model.score_response import LLMScoreResponse
 
-# data structures
+# data structure
 from src.search_engine.data_store import DataStore
 
 # build factories
@@ -17,70 +23,86 @@
 
 # logging
 from src.logger import configure_logging
-import logging
-
-
-if __name__ == "__main__":
-    args = parse_args()
+from logging import Logger, getLogger, DEBUG, INFO
 
-    config = Config.load(args.config_file)
 
-    if args.verbose:
-        configure_logging(logging.DEBUG)
+def get_and_setup_logging(verbose: bool = False) -> Logger:
+    if verbose:
+        configure_logging(DEBUG)
     else:
-        configure_logging(logging.INFO)
-    log = logging.getLogger(__name__)
+        configure_logging(INFO)
 
-    search_engine: BaseSearchEngine = SearchEngineFactory.build(search_engine_type=config.search_engine_type,
-                                              endpoint=config.search_engine_collection_endpoint)
-    data_store = DataStore()
+    return getLogger(__name__)
 
-    num_queries = 0
+def add_user_queries(config: Config, data_store: DataStore):
     if config.queries is not None:
         with open(config.queries, 'r', encoding='utf-8') as file:
             for line in file:
                 if line.strip():
                     data_store.add_query(line)
-                    num_queries += 1
 
-    # retrieval of the documents needed to generate the queries
-    docs_to_generate_queries = search_engine.fetch_for_query_generation(documents_filter=config.documents_filter,
-                                                                        doc_number=config.doc_number,
-                                                                        doc_fields=config.doc_fields)
-    log.debug(f"Number of documents retrieved for generation: {len(docs_to_generate_queries)}")
-    llm: BaseChatModel = LLMServiceFactory.build(LLMConfig.load(config.llm_configuration_file))
-    service = LLMService(chat_model=llm)
+def generate_and_add_queries(llm_service: LLMService, config: Config, data_store: DataStore) -> None:
+    num_queries_per_doc: int = int(((config.num_queries_needed - len(data_store.get_queries())) // config.doc_number) * 1.5)
 
-    num_queries_per_doc = int(( (config.num_queries_needed - num_queries) // config.doc_number) * 1.5)
-
-    # query generation step
-    all_queries_generated = False
     for doc in docs_to_generate_queries:
         data_store.add_document(doc.id, doc)
-        queries = service.generate_queries(doc, num_queries_per_doc)
-        for query_text in queries:
+        query_response: LLMQueryResponse = llm_service.generate_queries(doc, num_queries_per_doc)
+        for query_text in query_response.get_queries():
             if len(data_store.get_queries()) >= config.num_queries_needed:
-                all_queries_generated = True
-                break
-            query_id = data_store.add_query(query_text, doc.id)
+                return
+            query_id: str = data_store.add_query(query_text, doc.id)
             data_store.add_rating_score(query_id, doc.id, max(config.relevance_label_set))
-        if all_queries_generated:
-            break
 
-    log.debug(f"Number of documents evaluated: {len(docs_to_generate_queries)}")
 
-    # loop looking at all docs not rated in the data_store for that query
+def retrieve_and_add_documents(config: Config, data_store: DataStore) -> None:
+    for query_rating_context in data_store.get_queries():
+        docs_eval: List[Document] = search_engine.fetch_for_evaluation(keyword=query_rating_context.get_query(),
+                                                                       query_template=config.query_template,
+                                                                       doc_fields=config.doc_fields)
+        for doc in docs_eval:
+            if not data_store.has_document(doc.id):
+                data_store.add_document(doc.id, doc)
+
+def add_cartesian_product_scores(llm_service: LLMService, config: Config, data_store: DataStore) -> None:
     for query_rating_context in data_store.get_queries():
         for doc in data_store.get_documents():
             if not data_store.has_rating_score(query_rating_context.get_query_id(), doc.id):
-                score = service.generate_score(data_store.get_document(doc.id),
-                                               query_rating_context.get_query(),
-                                               config.relevance_scale)
+                score_response: LLMScoreResponse = llm_service.generate_score(data_store.get_document(doc.id),
+                                                        query_rating_context.get_query(),
+                                                        config.relevance_scale)
                 data_store.add_rating_score(query_rating_context.get_query_id(),
                                             doc.id,
-                                            score)
+                                            score_response.get_score())
+
+
+if __name__ == "__main__":
+    # configuration and logger definition
+    args = parse_args()
+    config: Config = Config.load(args.config_file)
+    log: Logger = get_and_setup_logging(args.verbose)
 
+    # setup
+    data_store: DataStore = DataStore()
+    search_engine: BaseSearchEngine = SearchEngineFactory.build(search_engine_type=config.search_engine_type,
+                                                                endpoint=config.search_engine_collection_endpoint)
+    llm: BaseChatModel = LLMServiceFactory.build(LLMConfig.load(config.llm_configuration_file))
+    service: LLMService = LLMService(chat_model=llm)
     writer: AbstractWriter = WriterFactory.build(config.output_format, data_store)
+
+    # pipeline starts
+    add_user_queries(config, data_store)
+
+    docs_to_generate_queries: List[Document] = search_engine.fetch_for_query_generation(documents_filter=config.documents_filter,
+                                                                                        doc_number=config.doc_number,
+                                                                                        doc_fields=config.doc_fields)
+    log.debug(f"Number of documents retrieved for generation: {len(docs_to_generate_queries)}")
+
+    generate_and_add_queries(service, config, data_store)
+
+    retrieve_and_add_documents(config, data_store)
+
+    add_cartesian_product_scores(service, config, data_store)
+
     writer.write(config.output_destination)
 
     log.info(f"Synthetic Dataset has been generated in: {config.output_destination}")

rre-dataset-generator/src/llm/llm_service.py

@@ -1,13 +1,13 @@
+import json
 from json import JSONDecodeError
-from typing import List, Union
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
-import json
+from src.model.query_response import LLMQueryResponse
+from src.model.score_response import LLMScoreResponse
 import logging
 
 from src.model.document import Document
-from typing import Union, List
 
 log = logging.getLogger(__name__)
 
@@ -16,7 +16,7 @@ class LLMService:
     def __init__(self, chat_model: BaseChatModel):
         self.chat_model = chat_model
 
-    def generate_queries(self, document: Document, num_queries_generate_per_doc: int) -> Union[List[str], str]:
+    def generate_queries(self, document: Document, num_queries_generate_per_doc: int) -> LLMQueryResponse:
         """
         Generate queries based on the given document and num_queries_generate_per_doc and
         Returns a list of generated queries or just a generated string in case of LLM hallucination
@@ -35,32 +35,35 @@ def generate_queries(self, document: Document, num_queries_generate_per_doc: int
             HumanMessage(content=f"Document:\n{doc_json}")
         ]
 
-        raw = self.chat_model.invoke(messages).content.strip()
+        # The response from invoke is an AIMessage object which contains all the needed info
+        response = self.chat_model.invoke(messages)
 
-        # in case LLM hallucinates
         try:
-            return json.loads(raw)
-        except JSONDecodeError:
-            log.warning("LLM hallucinated and its response: %s", raw)
-            return raw
+           output = LLMQueryResponse(response_content=response.content)
+        except (KeyError, JSONDecodeError, ValueError) as e:
+            log.warning(f"LLM unexpected response. Raw output: {response.content}")
+            raise ValueError(f"Invalid LLM response: {e}")
 
-    def generate_score(self, document: Document, query: str, relevance_scale: str) -> int:
+        return output
+    
+
+    def generate_score(self, document: Document, query: str, relevance_scale: str) -> LLMScoreResponse:
         """
         Generates a relevance score for a given document-query pair using a specified relevance scale.
         """
         if relevance_scale == "binary":
-            scale = {0, 1}
-            description = (" - 0: the query is NOT relevant to the given document"
-                           " - 1: the query is relevant to the given document")
+            allowed = {0, 1}
+            description = (" - 0: the query is NOT relevant to the given document\n"
+                        " - 1: the query is relevant to the given document")
         elif relevance_scale == "graded":
-            scale = {0, 1, 2}
-            description = (" - 0: the query is NOT relevant to the given document"
-                           " - 1: the query may be relevant to the given document"
-                           " - 2: the document proposed is the answer to the query")
+            allowed = {0, 1, 2}
+            description = (" - 0: the query is NOT relevant to the given document\n"
+                        " - 1: the query may be relevant to the given document\n"
+                        " - 2: the document proposed is the answer to the query")
         else:
-            error_msg = "The relevance scale must be either 'binary' or 'graded'"
-            log.error(error_msg)
-            raise ValueError(error_msg)
+            msg = f"Invalid relevance scale: {relevance_scale}"
+            log.error(msg)
+            raise ValueError(msg)
 
         messages = [
             SystemMessage(
@@ -76,16 +79,19 @@ def generate_score(self, document: Document, query: str, relevance_scale: str) -
             )
         ]
 
-        #response = self.chat_model.with_structured_output(method="json_mode").invoke(messages)
         raw = self.chat_model.invoke(messages).content.strip()
+
         try:
-            response =  int(json.loads(raw)['score'])
-            if response not in scale:
-                error_msg = f"LLM hallucinated the value of the scale. Returned: {response}"
-                log.warning(error_msg)
-                raise ValueError(error_msg)
-            return response
-        except JSONDecodeError:
-            error_msg = f"LLM hallucinated and its response: {raw}"
-            log.warning(error_msg)
-            raise ValueError(error_msg)
+            score = json.loads(raw)['score']
+        except (JSONDecodeError, KeyError) as e:
+            log.debug(f"LLM unexpected response. Raw output: {raw}")
+            raise ValueError(f"Invalid LLM response: {e}")
+
+        try:
+            parsed = LLMScoreResponse(score=score, scale=relevance_scale)
+            return parsed
+        except ValueError as e:
+            log.warning(f"Validation error for score '{score}' on scale '{relevance_scale}': {e}")
+            raise e
+
+

rre-dataset-generator/src/model/query_response.py

@@ -0,0 +1,40 @@
+import json
+from typing import List
+
+class LLMQueryResponse:
+    """
+    Parses and validates a JSON response containing a list of queries.
+    """
+    _queries: List[str]
+
+    def __init__(self, response_content: str):
+        """
+        Initializes the object by parsing and validating the JSON string.
+
+        Args:
+            response_content: A string containing a JSON-encoded list of strings.
+
+        Raises:
+            ValueError: If the content is not a valid JSON list of non-empty strings.
+        """
+        try:
+            parsed_content = json.loads(response_content)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in `response_content`: {e}")
+
+        if not isinstance(parsed_content, list):
+            raise ValueError("`response_content` must be a JSON list.")
+
+        if not all(isinstance(item, str) for item in parsed_content):
+            raise ValueError("All items in `response_content` must be strings.")
+
+        if any(not item.strip() for item in parsed_content):
+            raise ValueError("Items in `response_content` must not be empty or only whitespace.")
+
+        self._queries = parsed_content
+
+    def get_queries(self) -> List[str]:
+        """
+        Returns the validated list of queries.
+        """
+        return self._queries

rre-dataset-generator/src/model/score_response.py

@@ -0,0 +1,33 @@
+class LLMScoreResponse:
+    """
+    Parses and validates an LLM score response.
+    """
+    def __init__(self, score: int, scale: str = "graded"):
+        """
+        Initializes the object by validating the score.
+
+        Args:
+            score: The relevance score.
+            scale: The relevance scale, either 'binary' (0-1) or 'graded' (0-2).
+
+        Raises:
+            ValueError: If the score is not valid for the given scale.
+        """
+        if scale not in ["binary", "graded"]:
+            raise ValueError(f"Invalid scale: {scale}. Must be 'binary' or 'graded'.")
+            
+        if scale == "binary" and score not in {0, 1}:
+            raise ValueError(f"Score must be 0 or 1 for binary scale, got {score}")
+        elif scale == "graded" and score not in {0, 1, 2}:
+            raise ValueError(f"Score must be 0, 1, or 2 for graded scale, got {score}")
+            
+        self.score = score
+
+    def get_score(self) -> int:
+        """
+        Returns the validated score.
+
+        Returns:
+            The validated score.
+        """
+        return self.score