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:

1. Synthetic queries: Generated from chunks of text extracted from French administrative documents.
2. Retrieved documents: For each synthetic query, relevant documents were retrieved using the BM25 ranking algorithm.
3. 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 BGE for the embeddings and lancedb for the retrieval part. You can use your preffered embedding model to create the embeddings and add them to the database, LanceDB 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
import lancedb
from lancedb.pydantic import LanceModel, Vector
from lancedb.embeddings import get_registry

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

# 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."
    }
]

#BGE embedding model loading and creating the structure adapted to lance
model = get_registry().get("sentence-transformers").create(name="BAAI/bge-m3", device="cuda")

class Documents(LanceModel):
    main_title: str
    text: str = model.SourceField()
    hash: str
    vector: Vector(model.ndims()) = model.VectorField()

#Create table
table = db.create_table("example", schema=Documents, mode="overwrite")

def process_batch(batch):
    processed_documents = []
    for item in batch:
        try:
            processed_documents.append({
                "hash": item.get("hash", ""),
                "main_title": item.get("main_title", ""),
                "text": item.get("text", "")
                # Add any other fields you want to include
            })
        except Exception as e:
            print(f"Error processing item: {item}")
            print(f"Error message: {str(e)}")
    return processed_documents

# Process and add documents in batches
batch_size = 2  # Adjust as needed
for i in tqdm(range(0, len(documents), batch_size)):
    batch = documents[i:i+batch_size]
    processed_batch = process_batch(batch)
    if processed_batch:  # Only add if the batch is not empty
        table.add(processed_batch)

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