Update README.md

README.md

@@ -1,6 +1,7 @@
 # Cassandre-RAG
 
-Cassandre-RAG is a fine-tuned llama-3.1-8b model, built for RAG on French administrative documents, with a focus on sources from school administration.
+Cassandre-RAG is a fine-tuned **llama-3.1-8b model**, built for RAG on French administrative documents, with a focus on sources from school administration.
+The model has been trained to expect a predifined input structure, that allows it to very efficiently perform RAG tasks while clearly citing the specific exceprts and the source documents used in the generation of anwsers. 
 
 ## Training
 
@@ -46,11 +47,57 @@ Quantization:
 Cassandre-RAG uses a custom syntax for parsing sources and generating sourced output. 
 Each source should be preceded by an ID encapsulated in double asterisks (e.g., \*\*SOURCE_ID\*\*).
 
+The input structure expected by Cassandre is the following
+
+```yaml
+prompt = f"""### Query ###\n{user_message}\n\n### Source ###\n{fiches}\n\n### Answer ###\n"""
+```
+
+The **"Query"** section consists on the question or keywords that the user inputs.
+
+The **"Source"** part consists of the documents that have been retrieved using a vector database, like duckdb, lancedb or others.
+
+The **"Answer"** indicates the model where it should insert the generated answer to the query, based on the retrieved documents. 
+This answer will also contain the excerpts of the documents used and the ID of those documents, using this format:
+
+```yaml
+<ref text="[Quoted text from source]">[Source ID]</ref>
+```
+
 ### Example Usage
 
+In this example, we will be using lancedb for the retrieval part. Lancedb creates the embeddings of the documents added to the database, and its hybrid search feature allows us to combine vector search with keyword search for better retrieval.
+
 ```python
-import pandas as pd
+import lancedb
 from vllm import LLM, SamplingParams
+import pandas as pd
+
+# Initialize LanceDB
+db = lancedb.connect("lancedb_data")
+table = db.open_table("education")
+
+# Create fictitious education documents
+documents = [
+    {
+        "hash": "DOC001",
+        "main_title": "Inscription à l'école primaire",
+        "text": "L'inscription à l'école primaire en France se fait en deux étapes. Premièrement, les parents doivent se rendre à la mairie avec un justificatif de domicile, le livret de famille et le carnet de santé de l'enfant. Ensuite, ils doivent finaliser l'inscription directement à l'école. L'âge minimal pour l'inscription est de 3 ans."
+    },
+    {
+        "hash": "DOC002",
+        "main_title": "Calendrier des inscriptions scolaires",
+        "text": "Les inscriptions à l'école primaire doivent être effectuées au plus tard au mois de juin précédant la rentrée scolaire. Il est conseillé de s'y prendre à l'avance car certaines communes ont des périodes d'inscription spécifiques. La rentrée scolaire a généralement lieu début septembre."
+    },
+    {
+        "hash": "DOC003",
+        "main_title": "Documents requis pour l'inscription scolaire",
+        "text": "Pour inscrire un enfant à l'école primaire, les documents suivants sont généralement requis : justificatif de domicile de moins de 3 mois, livret de famille ou extrait d'acte de naissance, carnet de santé avec vaccinations à jour, et éventuellement le certificat de radiation si l'enfant était précédemment inscrit dans une autre école."
+    }
+]
+
+# Add documents to LanceDB (in a real scenario, this would be done separately)
+table.add(documents)
 
 # Load the model
 model_name = "PleIAs/Cassandre-RAG"
@@ -65,17 +112,24 @@ sampling_params = SamplingParams(
     stop=["#END#"]
 )
 
-# Prepare the input data
+def hybrid_search(text):
+    results = table.search(text, query_type="hybrid").limit(3).to_pandas()
+    document = []
+    for _, row in results.iterrows():
+        hash_id = str(row['hash'])
+        title = row['main_title']
+        content = row['text']
+        document.append(f"**{hash_id}**\n{title}\n{content}")
+    return "\n\n".join(document)
+
 def prepare_prompt(query, sources):
-    sources_text = "\n\n".join([f"**{src_id}**\n{content}" for src_id, content in sources])
-    return f"### Query ###\n{query}\n\n### Source ###\n{sources_text}\n\n### Analysis ###\n"
-
-# Example query and sources
-query = "Quelles sont les procédures pour inscrire un enfant à l'école primaire?"
-sources = [
-    ("SOURCE_001", "L'inscription à l'école primaire se fait généralement à la mairie..."),
-    ("SOURCE_002", "Les documents nécessaires pour l'inscription scolaire incluent..."),
-]
+    return f"### Query ###\n{query}\n\n### Source ###\n{sources}\n\n### Answer ###\n"
+
+# Example query
+query = "Quelles sont les démarches pour inscrire un enfant à l'école primaire en France?"
+
+# Perform hybrid search
+sources = hybrid_search(query)
 
 # Prepare the prompt
 prompt = prepare_prompt(query, sources)
@@ -85,5 +139,23 @@ outputs = llm.generate([prompt], sampling_params)
 generated_text = outputs[0].outputs[0].text
 
 print("Query:", query)
+print("\nSources:")
+print(sources)
 print("\nGenerated Response:")
-print(generated_text)
+print(generated_text)
+
+# Simple post-processing to display references (a simplified version of format_references)
+def simple_format_references(text):
+    while '<ref text="' in text:
+        start = text.index('<ref text="') + 10
+        end = text.index('">', start)
+        ref_text = text[start:end]
+        ref_id_start = end + 2
+        ref_id_end = text.index('</ref>', ref_id_start)
+        ref_id = text[ref_id_start:ref_id_end]
+        text = text.replace(f'<ref text="{ref_text}">{ref_id}</ref>', f'[{ref_id}: {ref_text}]')
+    return text
+
+print("\nFormatted Response:")
+print(simple_format_references(generated_text))
+```