app.py

import os
from pathlib import Path
from tempfile import mkdtemp

import gradio as gr
from docling_haystack.converter import ExportType
from dotenv import load_dotenv

from docling_haystack.converter import DoclingConverter
from haystack import Pipeline
from haystack.components.embedders import (
    SentenceTransformersDocumentEmbedder,
    SentenceTransformersTextEmbedder,
)
from haystack.components.preprocessors import DocumentSplitter
from haystack.components.writers import DocumentWriter
from milvus_haystack import MilvusDocumentStore, MilvusEmbeddingRetriever
from docling.chunking import DocChunk, HybridChunker
from haystack.components.builders import AnswerBuilder
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.components.generators import HuggingFaceAPIGenerator
from haystack.utils import Secret

def _get_env_from_colab_or_os(key):
    try:
        from google.colab import userdata

        try:
            return userdata.get(key)
        except userdata.SecretNotFoundError:
            pass
    except ImportError:
        pass
    return os.getenv(key)

load_dotenv()
HF_TOKEN = _get_env_from_colab_or_os("HF_TOKEN")
PATHS = ["Codex_over_het_welzijn_op_het_werk.pdf"]  # Docling Technical Report
EMBED_MODEL_ID = "intfloat/multilingual-e5-base" # sentence-transformers/all-MiniLM-L6-v2"
EXPORT_TYPE = ExportType.DOC_CHUNKS
TOP_K = 3
MILVUS_URI = str("codex-e5-base.db")

document_store = MilvusDocumentStore(
    connection_args={"uri": MILVUS_URI},
    drop_old=False,
    text_field="txt",  # set for preventing conflict with same-name metadata field
)

prompt_template = """
    Beantwoord de vraag met behulp van volgende documenten.
    Documenten:
    {% for doc in documents %}
        {{ doc.content }}
    {% endfor %}
    Vraag: {{query}}
    Antwoord:
    """

def create_rag_pipeline(model_id):
    rag_pipe = Pipeline()
    rag_pipe.add_component(
        "embedder",
        SentenceTransformersTextEmbedder(model=EMBED_MODEL_ID),
    )
    rag_pipe.add_component(
        "retriever",
        MilvusEmbeddingRetriever(document_store=document_store, top_k=TOP_K),
    )
    rag_pipe.add_component("prompt_builder", PromptBuilder(template=prompt_template))
    rag_pipe.add_component(
        "llm",
        HuggingFaceAPIGenerator(
            api_type="serverless_inference_api",
            api_params={"model": model_id},
            token=Secret.from_token(HF_TOKEN) if HF_TOKEN else None,
        ),
    )
    rag_pipe.add_component("answer_builder", AnswerBuilder())
    rag_pipe.connect("embedder.embedding", "retriever")
    rag_pipe.connect("retriever", "prompt_builder.documents")
    rag_pipe.connect("prompt_builder", "llm")
    rag_pipe.connect("llm.replies", "answer_builder.replies")
    rag_pipe.connect("llm.meta", "answer_builder.meta")
    rag_pipe.connect("retriever", "answer_builder.documents")
    return rag_pipe

def answer_question(question, selected_model_id):
    rag_pipe = create_rag_pipeline(selected_model_id)
    rag_res = rag_pipe.run(
        {
            "embedder": {"text": question},
            "prompt_builder": {"query": question},
            "answer_builder": {"query": question},
        }
    )

    answer = rag_res['answer_builder']['answers'][0].data.strip()
    sources = rag_res["answer_builder"]["answers"][0].documents

    sources_info = []
    for source in sources:
        if EXPORT_TYPE == ExportType.DOC_CHUNKS:
            doc_chunk = DocChunk.model_validate(source.meta["dl_meta"])
            sources_info.append(f"*** text: {doc_chunk.text}")
            if doc_chunk.meta.origin:
                sources_info.append(f"    file: {doc_chunk.meta.origin.filename}")
            if doc_chunk.meta.headings:
                sources_info.append(f"    section: {' / '.join(doc_chunk.meta.headings)}")
            bbox = doc_chunk.meta.doc_items[0].prov[0].bbox
            sources_info.append(
                f"    page: {doc_chunk.meta.doc_items[0].prov[0].page_no}, "
                f"    bounding box: [{int(bbox.l)}, {int(bbox.t)}, {int(bbox.r)}, {int(bbox.b)}]"
            )
        elif EXPORT_TYPE == ExportType.MARKDOWN:
            sources_info.append(repr(source.content))
        else:
            raise ValueError(f"Unexpected export type: {EXPORT_TYPE}")

    return answer, str("\n".join(sources_info))

def main():
    with gr.Blocks() as demo:
        gr.Markdown('''<sub><sup>(english explanation below)</sup></sub>

**RAG met codex welzijn op het werk**

Deze app biedt een slimme manier om specifieke vragen te stellen over de  **[Codex over het welzijn op het werk](https://huggingface.co/spaces/to-be/chat_met_codex_over_het_welzijn_op_het_werk/resolve/main/Codex_over_het_welzijn_op_het_werk.pdf)**, een Belgische wetgeving die de veiligheid, gezondheid en het welzijn van werknemers regelt. De codex bestaat uit 10 boeken die diverse aspecten van arbeidsomstandigheden behandelen, zoals preventie van psychosociale risico's, gezondheidstoezicht en eerste hulp. Dit kan een hulp zijn voor preventieadviseurs en werkpleksveiligheidsdeskundigen.

''')
        
        with gr.Row():
            question_input = gr.Textbox(label="Vraag", value="Mag een werkgever zelf een asbestinventaris maken voor zijn bedrijf?")
        with gr.Row():
            sample_questions = gr.Dropdown(
                label="Voorbeeldvragen",
                choices=[
                    "Mag een werkgever zelf een asbestinventaris maken voor zijn bedrijf?",
                    "Rangschik volgende beschermingsmaatregelen volgens prioriteit:\n A. Organisatorische maatregelen: opleidingen, procedures, werkvergunningen, …\n B. Risico’s verminderen met collectieve beschermingsmiddelen zoals relingen\n C. Substitutie van het gevaar: bv. een gevaarlijk chemisch product vervangen door een minder gevaarlijke stof\n D. Persoonlijke beschermingsmiddelen: denk aan een veiligheidsharnas",
                    "Wat is geen taak van de Interne Dienst voor Preventie en Bescherming op het Werk (IDPBW)? Kies 1 van volgende opties:\nA: Risico's onderzoeken en advies geven over de risico-evaluatie\nB: Advies geven over arbeidshygiëne\nC: Werk-privébalans voor werknemers bewaken\nD: Policies rond telewerken goedkeuren\nE: Werkinstructies opstellen"
                ],
                value="Mag een werkgever zelf een asbestinventaris maken voor zijn bedrijf?"
            )
        with gr.Row(equal_height=True):
            model_dropdown = gr.Dropdown(
                label="Selecteer generatie model",
                choices=[
                    "01-ai/Yi-1.5-34B-Chat",
                    "mistralai/Mixtral-8x7B-Instruct-v0.1",
                    "mistralai/Mistral-Nemo-Instruct-2407",
                    "Qwen/Qwen2.5-72B-Instruct",
                    "microsoft/Phi-3.5-mini-instruct"
                ],
                value="microsoft/Phi-3.5-mini-instruct"
            )
            send_button = gr.Button("Send")
        with gr.Row():
            answer_output = gr.Textbox(label="Antwoord")
            sources_output = gr.Textbox(label="Bronnen")

        gr.Markdown("<sub><sup>LLM kan fouten maken. Dubbelcheck belangrijke informatie</sup></sub>")
        gr.Markdown('''


**Hoe werkt de app?**

De app maakt gebruik van  **Retrieval-Augmented Generation (RAG)**  om specifieke vragen te beantwoorden op basis van de inhoud van de codex. De gebruiker kan een vraag stellen via een eenvoudige interface en krijgt een antwoord met **relevante bronnen** uit de codex.


**Technologieën die zijn gebruikt:**

-   **Docling**: Voor het converteren en indexeren van de documenten.
-   **Haystack**: Voor het bouwen van de RAG-pipeline.
-   **Hugging Face Inference API**: Voor het genereren van antwoorden via AI.
-   **Gradio**: Voor het bouwen van de gebruikersinterface.

**Verbeteringen:**

-   Het embedding-model is gewijzigd van  `sentence-transformers/all-MiniLM-L6-v2`  naar  `intfloat/multilingual-e5-base`  vanwege betere prestaties in multilinguïsme, vooral voor het Nederlands.
-   Je hebt keuze tussen verschillende LLM's, met persoonlijke voorkeur voor  **Phi-3.5-mini-instruct.** 
        
This demo provides a smart way to ask specific questions about the  **Codex on Well-Being at Work**, a Belgian legislation that regulates the safety, health, and well-being of employees. The codex consists of 10 books that cover various aspects of working conditions, such as prevention of psychosocial risks, health monitoring, and first aid. This can be a help for prevention advisors and workplace safety professionals.

**How does the app work?**

The app uses  **Retrieval-Augmented Generation (RAG)**  to answer specific questions based on the content of the codex. The user can ask a question through a simple interface and receives an answer with  **relevant sources**  from the codex.

**Technologies used:**

-   **Docling**: For converting and indexing the documents.
-   **Haystack**: For building the RAG pipeline.
-   **Hugging Face Inference API**: For generating answers via AI.
-   **Gradio**: For building the user interface.

**Improvements:**

-   The embedding model was changed from  `sentence-transformers/all-MiniLM-L6-v2`  to  `intfloat/multilingual-e5-base`  for better performance in multilingualism, especially for Dutch.
-   You have the option to choose between different LLMs, with a personal preference for  **Phi-3.5-mini-instruct**.


Comments / questions can be directed to [toon@neontreebot.be](mailto:toon@neontreebot.be?subject=to-be/chat_met_codex_over_het_welzijn_op_het_werk)

''')
        gr.HTML('''<a href="https://visitorbadge.io/status?path=chat_met_codex_over_het_welzijn_op_het_werk"><img src="https://api.visitorbadge.io/api/combined?path=chat_met_codex_over_het_welzijn_op_het_werk&countColor=%23263759" /></a>''')
        sample_questions.change(lambda x: x, inputs=sample_questions, outputs=question_input)
        send_button.click(answer_question, inputs=[question_input, model_dropdown], outputs=[answer_output, sources_output])
        question_input.submit(answer_question, inputs=[question_input, model_dropdown], outputs=[answer_output, sources_output])

    demo.launch()

if __name__ == "__main__":
    main()