NaiveCode2Text/Final version

.gitignore

@@ -1,5 +1,9 @@
 # Personnal usage
 test.ipynb
+test.py
+
+# Hydra
+outputs/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.13"
 dependencies = [
     "duckdb>=1.4.4",
     "fastparquet>=2025.12.0",
+    "hydra-core>=1.3.2",
     "ipykernel>=7.1.0",
     "ipywidgets>=8.1.8",
     "jupyter-cache>=1.0.1",

src/agents/NaiveCode2Text/code_retrieval/code_sampler.py

@@ -1,41 +1,13 @@
 import polars as pl
 import s3fs
-import numpy as np
-
-
-def sample_codes(
-        fs: s3fs.S3FileSystem,
-        population_path: str,
-        code_column: str,
-        n_codes: int
-        ) -> np.ndarray:
-    """
-    Sample codes with replacement using dataframes from Polars.
-
-    Args:
-        fs (S3FileSystem): The filesystem for importation.
-        population_path (str): The path of the parquet file of the population.
-        code_column (str): The name of the column for codes.
-        n_codes (int): The number of codes to sample.
-
-    Returns:
-        numpy.ndarray: An array of n_codes codes sampled with replacement.
-    """
-
-    with fs.open(population_path, 'rb') as f:
-        df = pl.read_parquet(f)
-
-    sampled = df.select(code_column).sample(n=n_codes, with_replacement=True)
-
-    return sampled[code_column].to_numpy()
 
 
 def sample_codes_lazy(
         fs: s3fs.S3FileSystem,
         population_path: str,
         code_column: str,
-        n_codes: str
-        ) -> np.ndarray:
+        n_codes: int
+        ) -> list:
     """
     Sample codes with replacement using lazyframes from Polars.
 
@@ -46,7 +18,7 @@ def sample_codes_lazy(
         n_codes (int): The number of codes to sample.
 
     Returns:
-        numpy.ndarray: An array of n_codes codes sampled with replacement.
+        list: A list of n_codes codes sampled with replacement.
     """
 
     with fs.open(population_path, 'rb') as f:
@@ -68,4 +40,4 @@ def sample_codes_lazy(
 
     df = sampled.collect()
 
-    return df[code_column].to_numpy()
+    return df[code_column].to_list()

src/agents/NaiveCode2Text/code_retrieval/code_specifier.py

@@ -1,8 +1,49 @@
-from src.neo4j_graph.graph import Graph
+import logging
+
+from langchain_neo4j import Neo4jGraph
+
+logger = logging.getLogger(__name__)
+
+
+def get_all_codes(
+        graph: Neo4jGraph,
+        n_samples_per_code: int
+        ) -> list:
+    """
+    Sample codes with replacement using lazyframes from Polars.
+
+    Args:
+        fs (S3FileSystem): The filesystem for importation.
+        population_path (str): The path of the parquet file of the population.
+        code_column (str): The name of the column for codes.
+        n_codes (int): The number of codes to sample.
+
+    Returns:
+        numpy.ndarray: An array of n_codes codes sampled with replacement.
+    """
+
+    query = """
+        MATCH (n)
+        WHERE n.FINAL = 1
+        RETURN DISTINCT n.CODE;
+        """
+
+    response = graph.query(query)
+
+    if not response:
+        logger.warn("No response in get_all_codes")
+        return []
+
+    result = []
+
+    for r in response:
+        result += [r["n.CODE"]] * n_samples_per_code
+
+    return result
 
 
 def get_code_information(
-        graph: Graph,
+        graph: Neo4jGraph,
         code: str
         ) -> dict:
     """
@@ -22,36 +63,81 @@ def get_code_information(
     OPTIONAL MATCH (node)-[:HAS_CHILD]->(child)
     WITH node, parent, collect({code: child.CODE, name: child.NAME}) as children
     RETURN node.CODE as code,
-    node.LEVEL as level,
-    node.NAME as name,
-    node.text as description,
-    node.Includes as includes,
-    node.IncludesAlso as includes_also,
-    node.Excludes as excludes,
-    node.Implementation_rule as implementation_rule,
-    parent.CODE as parent_code,
-    children,
-    size(children) as children_count
-    """
-    result = graph.graph.query(query, params={"code": code})
+           node.LEVEL as level,
+           node.NAME as name,
+           node.text as description,
+           node.Includes as includes,
+           node.IncludesAlso as includes_also,
+           node.Excludes as excludes,
+           node.Implementation_rule as implementation_rule,
+           parent.CODE as parent_code,
+           children,
+           size(children) as children_count
+    """
+    result = graph.query(query, params={"code": code})
 
     if not result:
-        print("No result in get_code_information")
-        return ()
+        logger.warn("No result in get_code_information")
+        return []
 
     return result[0]
 
 
-def NAF_to_NACE(
-        code: str
-        ) -> str:
+def get_code_list_information(
+        graph: Neo4jGraph,
+        code_list: list[str]
+        ) -> dict:
     """
-    For the case of NAF code format (DDDDL), transform it into NACE (DD.DD).
+    Retrieve code specifications from a Neo4j graph for multiple codes.
 
     Args:
-        code (str): The code in NAF format to transform.
+        graph (Neo4jGraph): The Neo4j graph.
+        code_list (list[str]): List of codes to retrieve.
 
     Returns:
-        str: The code in NACE format.
+        dict: Dictionary of code information keyed by code.
     """
-    return code[0:2] + '.' + code[2:4]
+
+    if not code_list:
+        return {}
+
+    query = """
+    MATCH (node)
+    WHERE node.CODE IN $code_list
+    OPTIONAL MATCH (node)<-[:HAS_CHILD]-(parent)
+    OPTIONAL MATCH (node)-[:HAS_CHILD]->(child)
+    WITH node, parent, collect({code: child.CODE, name: child.NAME}) AS children
+    RETURN node.CODE AS code,
+           node.LEVEL AS level,
+           node.NAME AS name,
+           node.text AS description,
+           node.Includes AS includes,
+           node.IncludesAlso AS includes_also,
+           node.Excludes AS excludes,
+           node.Implementation_rule AS implementation_rule,
+           parent.CODE AS parent_code,
+           children,
+           size(children) AS children_count
+    """
+
+    result = graph.query(query, params={"code_list": code_list})
+
+    if not result:
+        logger.warn("No result in get_code_list_information")
+        return {}
+
+    # Transformer la liste de résultats en dictionnaire de dictionnaires
+    code_dict = {}
+    for record in result:
+        code_key = record["code"]
+        code_dict[code_key] = {
+            "code": code_key,
+            "name": record["name"],
+            "description": record["description"],
+            "includes": record["includes"],
+            "includes_also": record["includes_also"],
+            "excludes": record["excludes"],
+            "implementation_rule": record["implementation_rule"]
+        }
+
+    return code_dict

src/agents/NaiveCode2Text/data_preprocessing/NAF_preprocessing.py

@@ -0,0 +1,49 @@
+"""
+Useful if the code of the input data for sampling is different from the one
+of the notice (DD.DDL).
+"""
+
+
+def NAF_to_NACE(
+        code: str
+        ) -> str:
+    """
+    For the case of NAF code format (DDDDL), transform it into NACE (DD.DD).
+
+    Args:
+        code (str): The code in NAF format to transform.
+
+    Returns:
+        str: The code in NACE format.
+    """
+    return code[0:2] + '.' + code[2:4]
+
+
+def to_proper_NAF(
+        code: str
+        ) -> str:
+    """
+    For the case of bad NAF code format (DDDDL), transform it into proper NAF (DD.DDL).
+
+    Args:
+        code (str): The code in bad NAF format to transform.
+
+    Returns:
+        str: The code in proper NAF format.
+    """
+    return code[0:2] + '.' + code[2:]
+
+
+def to_bad_NAF(
+        code: str
+        ) -> str:
+    """
+    For the case of proper NAF code format (DD.DDL), transform it into bad NAF (DDDDL).
+
+    Args:
+        code (str): The code in proper NAF format to transform.
+
+    Returns:
+        str: The code in bad NAF format.
+    """
+    return code[0:2] + code[3:]

src/agents/NaiveCode2Text/fewshot/fewshot_prompt_builder.py

@@ -0,0 +1,69 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def build_fewshot_system_prompt(
+        prompt_path: str,
+        language: str = "English",
+        nb_labels: int = 10,
+        ) -> str:
+    """
+    Import system prompt for fewshots with the correct language and specify the number of labels.
+
+    Args:
+        prompt_path (str): The path for importation.
+        language (str): English or French.
+        nb_labels (int): The number of labels to generate.
+
+    Returns:
+        str: The system prompt.
+    """
+    if language == "French":
+        suffix = "_fr"
+    elif language == "English":
+        suffix = "_en"
+    else:
+        logger.warn("Supported languages are English and French. Switching to English...")
+        suffix = "_en"
+
+    # Importing prompt
+    file_path = prompt_path + "system_prompt_fewshot" + suffix + ".txt"
+    with open(file=file_path, mode='r') as f:
+        system_prompt = f.read()
+
+    # Indicating the correct number of labels
+    system_prompt = system_prompt.replace("{nb_labels}", str(nb_labels))
+    return system_prompt
+
+
+def add_fewshot_user_prompt(
+        fewshot: list,
+        language: str = "English"
+        ) -> str:
+    """
+    Add few-shot to user prompt with the correct language.
+
+    Args:
+        fewshot (list): The list of exampls to give for few-shot.
+        language (str): English or French.
+
+    Returns:
+        str: The text to add to user prompt.
+    """
+
+    if len(fewshot) == 0:
+        return ""
+
+    if language == "English":
+        prompt = "\nHere are examples of labels for the code:"
+    elif language == "French":
+        prompt = "\nVoici des exemples de libellés pour le code:"
+    else:
+        logger.warn("Supported languages are English and French. Switching to English...")
+        prompt = "\nHere are examples of labels for the code:"
+
+    for example in fewshot:
+        prompt += "\n- " + example
+
+    return prompt