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. 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

The model was fine-tuned on a specialized corpus consisting of:

Synthetic queries: Generated from chunks of text extracted from French administrative documents.
Retrieved documents: For each synthetic query, relevant documents were retrieved using the BM25 ranking algorithm.
Generated answers: Responses to the synthetic queries were created based on the retrieved documents.

Training Hyperparameters:
  Max Steps: 3000
  Learning Rate: 3e-4
  Batch Size: 2 per device
  Gradient Accumulation Steps: 4
  Max Sequence Length: 8192
  Weight Decay: 0.001
  Warmup Ratio: 0.03
  LR Scheduler: Linear
  Optimizer: paged_adamw_32bit

LoRA Configuration:
  LoRA Alpha: 16
  LoRA Dropout: 0.1
  LoRA R: 64
  Target Modules: 
    - gate_proj
    - down_proj
    - up_proj
    - q_proj
    - v_proj
    - k_proj
    - o_proj

Quantization:
  Quantization: 4-bit
  Quantization Type: nf4
  Compute Dtype: float16

Usage

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

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:

<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.

import lancedb
from vllm import LLM, SamplingParams
import pandas as pd

# Initialize LanceDB
db = lancedb.connect("lancedb_data")
table = db.open_table("education")

# We will create some ficticious education documents to add to the database
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 
table.add(documents)

# Load the model
model_name = "PleIAs/Cassandre-RAG"
llm = LLM(model_name, max_model_len=8128)

# Set sampling parameters
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=3000,
    presence_penalty=1.2,
    stop=["#END#"]
)

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):
    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)

# Generate the response
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("\nFormatted Response:")
print(simple_format_references(generated_text))

Here we have an example of the respnse we will get:

### Query ###
Quelles sont les démarches pour inscrire un enfant à l'école primaire en France?

### Source ###
**DOC001**
Inscription à l'école primaire
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.

**DOC002**
Calendrier des inscriptions scolaires
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.

**DOC003**
Documents requis pour l'inscription scolaire
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.

### Answer ###
Pour inscrire un enfant à l'école primaire en France :

1. Allez à la mairie avec les documents nécessaires. <ref 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">DOC001</ref>

2. Inscrivez-vous avant juin pour la rentrée de septembre. <ref text="Les inscriptions à l'école primaire doivent être effectuées au plus tard au mois de juin précédant la rentrée scolaire">DOC002</ref>

3. L'enfant doit avoir au moins 3 ans. <ref text="L'âge minimal pour l'inscription est de 3 ans">DOC001</ref>

4. Finalisez l'inscription à l'école. <ref text="Ensuite, ils doivent finaliser l'inscription directement à l'école">DOC001</ref>

Apportez un certificat de radiation si l'enfant change d'école. <ref text="et éventuellement le certificat de radiation si l'enfant était précédemment inscrit dans une autre école">DOC003</ref>

Contactez votre mairie pour plus d'informations.

#END#