Merge branch 'master' into DAGE-18/Persistent-Storage-JSON

Author: MrCabss69

.github/workflows/test.yml

@@ -0,0 +1,31 @@
+name: Unit Tests
+
+on:
+  pull_request:
+    branches: [ main, dataset-generator ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.11"
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v3
+      with:
+        version: "latest"
+
+    - name: Install dependencies
+      run: uv sync
+      working-directory: rre-dataset-generator
+
+    - name: Run unit tests
+      run: uv run pytest
+      working-directory: rre-dataset-generator

rre-dataset-generator/config.yaml

@@ -2,28 +2,28 @@
 # dataset-generator module
 # Default: "q=#$query##"
 # Example (solr):
-query_template: "q=#$query##&fq=genre:horror&wt=json"
+query_template: "q=#$query##&df=description&rows=3&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
-search_engine_collection_endpoint: "http://localhost:8983/solr/mycore"
+search_engine_collection_endpoint: "http://localhost:8983/solr/testcore/"
 
 # (Optional) Filter query to restrict the set of documents used to generate queries.
 # If a field has more than one value, the documents retrieved have at least one of the values in the field (OR-
 # -like). If there are more than one fields, the documents are filtered for both fields (AND-like)
 # Default: all documents are considered and no filter is applied
-documents_filter:
-  - genre:
-      - "horror"
-      - "fantasy"
-  - type:
-      - "book"
+#documents_filter:
+#  - genre:
+#      - "horror"
+#      - "fantasy"
+#  - type:
+#      - "book"
 
 # Maximum number of documents to retrieve from the search engine to generate queries
-doc_number: 100
+doc_number: 2
 
 # List of fields from documents used to generate queries and relevance scoring
 # These should match fields available in the search engine schema
@@ -34,14 +34,14 @@ doc_fields:
 
 # (Optional) File containing predefined queries to use instead of or alongside generated ones
 # If available, must be in a txt format.
-queries: "queries.txt"
+#queries: "queries.txt"
 
 # (Optional) Whether to generate queries from documents
 # Default: true; set to false to disable query generation from documents
 generate_queries_from_documents: true
 
 # Total number of queries to generate (includes predefined queries, if any)
-num_queries_needed: 10
+num_queries_needed: 4
 
 # Relevance scale used for scoring document relevance
 # Accepted values:
@@ -57,4 +57,4 @@ llm_configuration_file: "llm_config.yaml"
 output_format: "quepid"
 
 # Path where the output dataset will be saved
-output_destination: "output/generated_dataset.json"
+output_destination: "output/generated_dataset.csv"

rre-dataset-generator/dataset_generator.py

@@ -1,27 +1,86 @@
-import logging
-
-from src.config import Config
+# configuration params
+from langchain_core.language_models import BaseChatModel
+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.search_engine.solr_search_engine import SolrSearchEngine
+# data structures
+from src.search_engine.data_store import DataStore
+
+# build factories
+from src.llm.llm_provider_factory import LLMServiceFactory
+from src.writers.writer_factory import WriterFactory
+from src.search_engine.search_engine_factory import SearchEngineFactory
+
+# logging
 from src.logger import configure_logging
 import logging
 
+
 if __name__ == "__main__":
     args = parse_args()
 
     config = Config.load(args.config_file)
+
     if args.verbose:
         configure_logging(logging.DEBUG)
     else:
         configure_logging(logging.INFO)
+    log = logging.getLogger(__name__)
+
+    search_engine: BaseSearchEngine = SearchEngineFactory.build(search_engine_type=config.search_engine_type,
+                                              endpoint=config.search_engine_collection_endpoint)
+    data_store = DataStore()
+
+    num_queries = 0
+    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)
+
+    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:
+            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)
+            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)}")
 
-    search_engine = SolrSearchEngine('http://localhost:8983/solr/testcore/')
+    # loop looking at all docs not rated in the data_store for that query
+    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)
+                data_store.add_rating_score(query_rating_context.get_query_id(),
+                                            doc.id,
+                                            score)
 
-    docs = search_engine.fetch_for_query_generation(documents_filter=config.documents_filter,
-                                                               doc_number=config.doc_number,
-                                                               doc_fields=config.doc_fields)
+    writer: AbstractWriter = WriterFactory.build(config.output_format, data_store)
+    writer.write(config.output_destination)
 
-    # docs = search_engine.fetch_for_evaluation(keyword="and",
-    #                                           query_template=config.query_template,
-    #                                           doc_fields=config.doc_fields)
+    log.info(f"Synthetic Dataset has been generated in: {config.output_destination}")

rre-dataset-generator/src/config.py

@@ -50,6 +50,20 @@ def check_config_type(cls, v):
                 raise ValueError("LLM_config file must have .yaml extension")
         return v
 
+    @property
+    def relevance_label_set(self) -> set[int]:
+        """
+        Returns the set of valid labels based on the relevance scale.
+        """
+        if self.relevance_scale == "binary":
+            return {0, 1}
+        elif self.relevance_scale == "graded":
+            return {0, 1, 2}
+        else:
+            error_msg = f"Unknown relevance scale: {self.relevance_scale}"
+            log.error(error_msg)
+            raise ValueError(error_msg)
+
     @classmethod
     def load(cls, config_path: str) -> Config:
         """

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

@@ -14,6 +14,8 @@
 from langchain_openai import ChatOpenAI
 from langchain_google_genai import ChatGoogleGenerativeAI
 from .llm_config import LLMConfig
+from langchain_core.language_models import BaseChatModel
+
 
 log = logging.getLogger(__name__)
 
@@ -45,17 +47,17 @@ def build_gemini(config: LLMConfig) -> BaseChatModel:
         google_api_key=key,
     )
 
-
-PROVIDER_REGISTRY = {
-    "openai": build_openai,
-    "gemini": build_gemini,
-}
-
-
-def build_chat_model(config: LLMConfig) -> BaseChatModel:
-    provider_name = config.name
-    if provider_name not in PROVIDER_REGISTRY:
-        log.error("Unsupported LLM provider requested: %s", provider_name)
-        raise ValueError(f"Unsupported provider: {provider_name}")
-    log.info("Selected LLM provider: %s", provider_name)
-    return PROVIDER_REGISTRY[provider_name](config)
+class LLMServiceFactory:
+    PROVIDER_REGISTRY = {
+        "openai": build_openai,
+        "gemini": build_gemini,
+    }
+
+    @classmethod
+    def build(cls, config: LLMConfig) -> BaseChatModel:
+        provider_name = config.name
+        if provider_name not in cls.PROVIDER_REGISTRY:
+            log.error("Unsupported LLM provider requested: %s", provider_name)
+            raise ValueError(f"Unsupported provider: {provider_name}")
+        log.info("Selected LLM provider: %s", provider_name)
+        return cls.PROVIDER_REGISTRY[provider_name](config)

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

@@ -1,12 +1,13 @@
 from json import JSONDecodeError
-from typing import Union, List
+from typing import List, Union
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
 import json
 import logging
 
 from src.model.document import Document
+from typing import Union, List
 
 log = logging.getLogger(__name__)
 
@@ -22,7 +23,7 @@ def generate_queries(self, document: Document, num_queries_generate_per_doc: int
         """
         system_prompt = (
             f"You are a helpful assistant! Generate {num_queries_generate_per_doc} "
-            "keyword-phrase-based queries based on the given document below. "
+            "queries based on the given document below. "
             "**Output only** a JSON array of strings—nothing else. "
             "Example format: [\"first query\", \"second query\"]"
         )

rre-dataset-generator/src/logger.py

@@ -20,7 +20,7 @@ def configure_logging(level: Union[str, int] = logging.INFO) -> None:
     logging.basicConfig(
         # filename='logs/log_file_name.log',
         level=level,
-        format='[%(asctime)s] %(filename)-15s:%(lineno)-4d %(levelname)-7s - %(message)s',
+        format='[%(asctime)s] %(filename)-28s:%(lineno)-4d %(levelname)-7s - %(message)s',
         datefmt='%H:%M:%S'
     )
 

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

@@ -17,7 +17,6 @@ class Document(BaseModel):
         ...,
         description="Fields of the document."
     )
-
     @field_validator('fields')
     @classmethod
     def check_no_empty_fields(cls, v: Dict[str, Any]) -> Dict[str, Any]:

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

@@ -40,7 +40,7 @@ def add_rating_score(self, doc_id: str, rating_score: int) -> None:
         self._doc_id_to_rating_score[doc_id] = rating_score
 
     def has_rating_score(self, doc_id: str) -> bool:
-        return self._doc_id_to_rating_score[doc_id] != self.DOC_NOT_RATED
+        return doc_id in self._doc_id_to_rating_score and self._doc_id_to_rating_score[doc_id] != self.DOC_NOT_RATED
 
     def get_rating_score(self, doc_id: str) -> int:
         return self._doc_id_to_rating_score[doc_id]

rre-dataset-generator/src/search_engine/data_store.py

@@ -5,11 +5,9 @@
 from pathlib import Path
 from typing import Dict, List, Optional
 
-from src.logger import configure_logging
 from src.model.document import Document
 from src.model.query_rating_context import QueryRatingContext
 
-configure_logging(level=logging.INFO)
 log = logging.getLogger(__name__)
 
 
@@ -64,7 +62,7 @@ def get_documents(self) -> List[Document]:
         """
         return list(self._documents.values())
 
-    def add_query(self, query: str, doc_id: str) -> str:
+    def add_query(self, query: str, doc_id: str | None = None) -> str:
         """
         If `query` already exists, just adds `doc_id` to it.
         Otherwise, creates a new QueryRatingContext.
@@ -73,7 +71,8 @@ def add_query(self, query: str, doc_id: str) -> str:
         if query in self._query_text_to_query_id:
             query_id = self._query_text_to_query_id[query]
             context = self._queries_by_id[query_id]
-            context.add_doc_id(doc_id)
+            if doc_id is not None:
+                context.add_doc_id(doc_id)
             return query_id
 
         # new query rating context