From 515b540873722277eadcfe02f237fdf93b93427a Mon Sep 17 00:00:00 2001
From: lgaliana <lino.galiana@insee.fr>
Date: Wed, 4 Dec 2024 13:21:03 +0000
Subject: [PATCH] ruff clean

---
 minimal-rag-in-notebook.ipynb         | 46 ++++++++++++++++++++++++---
 run_build_database.py                 |  6 +---
 src/model_building/build_llm_model.py | 14 ++++++--
 src/reranking/reranking_functions.py  | 32 +++++++++++--------
 src/utils/__init__.py                 | 13 ++++----
 src/utils/loading_utilities.py        |  7 ----
 6 files changed, 78 insertions(+), 40 deletions(-)

diff --git a/minimal-rag-in-notebook.ipynb b/minimal-rag-in-notebook.ipynb
index 121680e..3c728c9 100644
--- a/minimal-rag-in-notebook.ipynb
+++ b/minimal-rag-in-notebook.ipynb
@@ -320,15 +320,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Batches: 100%|██████████| 1/1 [00:00<00:00, 82.09it/s]\n",
-      "Processed prompts: 100%|██████████| 1/1 [00:01<00:00,  1.56s/it, est. speed input: 612.17 toks/s, output: 137.17 toks/s]\n"
+      "Batches: 100%|██████████| 1/1 [00:00<00:00, 88.49it/s]\n",
+      "Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Processed prompts: 100%|██████████| 1/1 [00:05<00:00,  5.53s/it, est. speed input: 172.93 toks/s, output: 136.89 toks/s]\n"
      ]
     }
    ],
@@ -358,7 +365,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
@@ -493,6 +500,37 @@
     "retriever.invoke(\"Quelle est la définition du PIB ?\")[:5]"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Batches: 100%|██████████| 1/1 [00:00<00:00, 75.98it/s]\n",
+      "Processed prompts: 100%|██████████| 1/1 [00:04<00:00,  4.60s/it, est. speed input: 707.17 toks/s, output: 131.72 toks/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[\" pour assister votre réponse, je vous fais un exposé succinct sur la définition du PIB et où prendre les chiffres les plus récents.\\n\\n     Le produit intérieur brut (PIB) est l'indicateur économique le plus important pour représenter la richesse monétaire générée par un appareil productif national ou régional. Il s'agit de la somme des valeurs ajoutées nouvellement créées par les unités productrices résidentes dans une année donnée, évaluées au prix du marché. Les chiffres du PIB peuvent être utiles pour faire des comparaisons internationales et jusqu'en 2018, ils sont publiés trois ans après le millésime.\\n\\n     La définition et les méthodes de calcul du PIB sont établies par l'Organisation des Nations Unies (ONU) et sont les mêmes pour tous les pays. L'application est contrôlée par l'office statistique de l'Union européenne (Eurostat). Il est important de savoir que le PIB actuel présente des inconvénients puisqu'il ne mesure pas efficacement le développement social ou la question de la soutenabilité de la croissance économique, notamment en termes écologiques. L'Insee présente donc de nouveaux outils afin de comprendre la mesure du PIB et de la croissance économique.\\n\\n     Pour trouver les chiffres les plus récents du PIB, vous pouvez consulter les statistiques de l'Insee sur la page [Comptes nationaux](https://www.insee.fr/fr/statistiques?taille=100&debut=0&theme=33&categorie=5&geo=FRANCE-1). Il existe deux types de publications, les comptes annuels et les comptes trimestriels. Les comptes annuels, publiés tous les ans, donnent une vision globale de la croissance économique pendant l'année et les comptes trimestriels, publiés chaque trimestre, fournissent une vision plus précise de l'évolution de la croissance économique au fil du temps.\\n\\n     Références :\\n\\n     - https://www.insee.fr/fr/information/2549709\\n     - https://www.insee.fr/fr/statistiques/2384072\\n     - https://www.insee.fr/fr/statistiques/4241527\"]"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "rag_chain.batch([\n",
+    "    \"Quelle est la définition du PIB ?\"\n",
+    "    \"Où trouver les nouveaux chiffres du chpimage ?\"\n",
+    "    ])"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/run_build_database.py b/run_build_database.py
index 2d746cc..c61a142 100644
--- a/run_build_database.py
+++ b/run_build_database.py
@@ -11,11 +11,7 @@
 import yaml
 
 from src.config import DefaultFullConfig, FullConfig, process_args, simple_argparser
-from src.db_building import (
-    build_or_load_document_database,
-    build_vector_database,
-    load_vector_database
-)
+from src.db_building import build_or_load_document_database, build_vector_database, load_vector_database
 
 # Logging configuration
 logger = logging.getLogger(__name__)
diff --git a/src/model_building/build_llm_model.py b/src/model_building/build_llm_model.py
index bcc0c0c..b8831a4 100644
--- a/src/model_building/build_llm_model.py
+++ b/src/model_building/build_llm_model.py
@@ -2,6 +2,7 @@
 import sys
 from typing import Any
 
+from langchain_huggingface.llms import HuggingFacePipeline
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
@@ -15,8 +16,8 @@
 
 from .fetch_llm_model import cache_model_from_hf_hub
 
-# Add the project root directory to sys.path
 root_dir = os.path.abspath(os.path.join(os.path.dirname(""), "./src"))
+
 if root_dir not in sys.path:
     sys.path.append(root_dir)
 
@@ -55,13 +56,20 @@ def build_llm_model(
         ),
         # Load LLM config
         "config": (
-            AutoConfig.from_pretrained(model_name, trust_remote_code=True, token=hf_token) if load_LLM_config else None
+            AutoConfig.from_pretrained(
+                model_name, trust_remote_code=True,
+                token=hf_token
+            ) if load_LLM_config else None
         ),
         "token": hf_token,
     }
 
     # Load LLM tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, device_map="auto", token=hf_token)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, use_fast=True,
+        device_map="auto", token=hf_token
+    )
+
     streamer = None
     if streaming:
         streamer = TextStreamer(tokenizer=tokenizer, skip_prompt=True)
diff --git a/src/reranking/reranking_functions.py b/src/reranking/reranking_functions.py
index 381ffeb..09a4ee3 100644
--- a/src/reranking/reranking_functions.py
+++ b/src/reranking/reranking_functions.py
@@ -18,9 +18,6 @@ def compress_documents_lambda(
     return retriever.get_relevant_documents(query)
 
 
-###### LLM Reraner functions ######
-
-
 def expected_relevance_values(logits, grades_token_ids, list_grades):
     next_token_logits = logits[:, -1, :]
     next_token_logits = next_token_logits.cpu()[0]
@@ -81,7 +78,9 @@ def RG_S(tokenizer, model, query, document, aggregating_method, k=5):
         },
     ]
 
-    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    input_text = tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, tokenize=False
+    )
 
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
 
@@ -117,9 +116,9 @@ def RG_4L(tokenizer, model, query, document, args):
 
     log_probs = []
     for judgement in possible_judgements:
-        input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False).format(
-            query=query, document=document, judgement=judgement
-        )
+        input_text = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=False, tokenize=False
+        ).format(query=query, document=document, judgement=judgement)
         log_probs.append(compute_sequence_log_probs(sequence=input_text))
 
     probs = F.softmax(torch.tensor(log_probs), dim=-1).numpy()
@@ -146,9 +145,9 @@ def RG_3L(tokenizer, model, query, document, args):
 
     log_probs = []
     for judgement in possible_judgements:
-        input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False).format(
-            query=query, document=document, judgement=judgement
-        )
+        input_text = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=False, tokenize=False
+        ).format(query=query, document=document, judgement=judgement)
         log_probs.append(compute_sequence_log_probs(sequence=input_text))
 
     probs = F.softmax(torch.tensor(log_probs), dim=-1).numpy()
@@ -177,7 +176,9 @@ def RG_YN(tokenizer, model, query, document, aggregating_method):
         },
     ]
 
-    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    input_text = tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, tokenize=False
+    )
 
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
 
@@ -188,8 +189,12 @@ def RG_YN(tokenizer, model, query, document, aggregating_method):
     return aggregating_method(logits, grades_token_ids, list_grades)
 
 
-def llm_reranking(tokenizer, model, query, retrieved_documents, assessing_method, aggregating_method):
-    docs_content = retrieved_documents.copy()  # [doc.page_content for doc in retrieved_documents]
+def llm_reranking(
+    tokenizer, model, query, retrieved_documents, assessing_method, aggregating_method
+):
+    docs_content = (
+        retrieved_documents.copy()
+    )  # [doc.page_content for doc in retrieved_documents]
 
     scores = []
     for document in docs_content:
@@ -200,4 +205,3 @@ def llm_reranking(tokenizer, model, query, retrieved_documents, assessing_method
     docs_with_scores.sort(key=lambda x: x[1], reverse=True)
     sorted_documents = [doc for doc, score in docs_with_scores]  # docs_with_scores
     return sorted_documents
-
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index d353fad..5077093 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -1,12 +1,11 @@
+from .formatting_utilities import create_prompt_from_instructions, format_docs
 from .loading_utilities import load_dataframe_from_parquet
 from .misc import compare_params
-from .formatting_utilities import (
-    create_prompt_from_instructions, format_docs
-)
 
 __all__ = [
-    "load_dataframe_from_parquet", "cache_model_from_hf_hub",
-    "compare_params", "create_prompt_from_instructions",
-    "format_docs"
+    "load_dataframe_from_parquet",
+    "cache_model_from_hf_hub",
+    "compare_params",
+    "create_prompt_from_instructions",
+    "format_docs",
 ]
-
diff --git a/src/utils/loading_utilities.py b/src/utils/loading_utilities.py
index 2581a25..ab91fcb 100644
--- a/src/utils/loading_utilities.py
+++ b/src/utils/loading_utilities.py
@@ -1,16 +1,9 @@
-import os
-import subprocess
 from typing import Literal
 
 import pandas as pd
-import s3fs
-
-from src.config import Configurable, DefaultFullConfig, FullConfig
 
 
 def load_dataframe_from_parquet(
     path: str, engine: Literal["auto", "pyarrow", "fastparquet"] = "fastparquet"
 ) -> pd.DataFrame:
     return pd.read_parquet(path, engine)
-
-