chatpdf.py

import os
import click
import fitz  # PyMuPDF
import numpy as np
import itertools as it
import operator as op
import requests
import tiktoken 
from typing import List, Tuple
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

load_dotenv()
protocol_buffers_python_implementation = os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"]
CACHE_FOLDER = os.environ["CACHE_FOLDER"]
CACHE_FOLDER = "./cache"  # Chemin du dossier cache

# Fonctions existantes pour convertir les PDFs et diviser les pages en chunks
def convert_pdf_to_text(pdf_data: str) -> List[str]:
    document = fitz.open(pdf_data)
    accumulator: List[str] = []
    for page_num in range(len(document)):
        page = document[page_num]
        text = page.get_text()
        accumulator.append(text)
    document.close()
    return accumulator

def split_pages_into_chunks(pages: List[str], chunk_size: int, tokenizer) -> List[str]:
    page_tokens: List[List[int]] = [tokenizer.encode(page) for page in pages]
    document_tokens = list(it.chain(*page_tokens))

    nb_tokens = len(document_tokens)
    nb_partitions = round(nb_tokens / chunk_size)

    accumulator: List[str] = []

    for chunck_tokens in np.array_split(document_tokens, nb_partitions):
        paragraph = tokenizer.decode(chunck_tokens)
        accumulator.append(paragraph)

    return accumulator

def vectorize(chunks: List[str], transformer: SentenceTransformer, device: str = 'cpu') -> List[Tuple[str, np.ndarray]]:
    embeddings = transformer.encode(
        sentences=chunks,
        batch_size=32,
        device=device
        ,show_progress_bar=True,
    )

    return list(zip(chunks, embeddings))

def find_candidates(query_embedding: np.ndarray, chunks: List[str], corpus_embeddings: np.ndarray, top_k: int = 5) -> List[str]:
    dot_product = query_embedding @ corpus_embeddings.T
    norms = np.linalg.norm(query_embedding) * np.linalg.norm(corpus_embeddings, axis=1)

    weighted_scores = dot_product / norms  # cosine similarity

    zipped_chunks_scores = list(zip(chunks, weighted_scores))
    sorted_chunks_scores = sorted(zipped_chunks_scores, key=op.itemgetter(1), reverse=True)
    selected_candidates = sorted_chunks_scores[:top_k]

    return list(map(op.itemgetter(0), selected_candidates))


def generate_prompt(question, chunks, corpus_embeddings, embedding_model, llm_model, top_k, document_type: str = "ce pdf"):
    query_embedding = embedding_model.encode([question])[0]
    candidates = find_candidates(query_embedding=query_embedding, 
                                 chunks=chunks, 
                                 corpus_embeddings=corpus_embeddings, 
                                 top_k=top_k)
    context = "\n".join(candidates)
    prompt = f"""
     ROLE: Tu es un assistant QA. Tu as été crée par Papa Séga de Orange INNVO. 
           Ton rôle est d'aider les utilisateurs à trouver la bonne réponse.
     FONCTIONNEMENT: 
            Voici un ensemble de documents dans {document_type}: {context}.
            Tu dois analyser ces documents pour répondre à la question de l'utilisateur. 
            Tu dois répondre aux utilisateurs uniquement en français.
            Tu ne dois pas répondre en anglais.
            Vérifie bien que la question est en rapport avec ces documents :
                -SI OUI ALORS :
                    - construit une réponse.
                    - tu peux reformuler la réponse
                - SI NON ALORS:
                    - Si la question est en rapport avec ton fonctionnement alors:
                         - rappelle-lui que tu es juste un modèle de QA et ne réponds pas à la question.
                    Si Non Alors: 
                            - il ne faut jamais répondre à cette question qui a été posée par l'utilisateur ! 
        RAPPEL :
            Ton objectif est d'aider l'utilisateur à trouver les réponses pertinentes sur ces questions !

    Question de l'utilisateur : {question}
    SI la question posée est hors contexte,
      - ALORS informe l'utilisateur que la question est hors contexte et que tu ne pourras lui donner une réponse.

    Attention 
      - Tu dois être courtois avec les utilisateurs 
      - Tu n'as pas le droit de faire des hallucinations.
    
    Réponds à la question en te basant UNIQUEMENT sur le contexte suivant :\n{context}\nQuestion : {question}

    """
    
    chain_prompt = ChatPromptTemplate.from_template(prompt)
    chain = chain_prompt | llm_model | StrOutputParser()

    response = chain.invoke({"topic": question})

    return response

@click.command()
@click.option('--pdf_url', required=True, help='URL du fichier PDF à télécharger')
@click.option('--pdf_path', default='document.pdf', help='Chemin local pour enregistrer le fichier PDF téléchargé')
@click.option('--embedding_model', required=True, help="Modèle d'embedding à utiliser")
@click.option('--llm_model', required=True, help='Modèle LLM à utiliser')
@click.option('--top_k', type=int, default=5, help='Nombre de candidats pour la recherche de similarité')
def main(pdf_url, pdf_path, embedding_model, llm_model, top_k):
    os.makedirs(CACHE_FOLDER, exist_ok=True)
    
    response = requests.get(pdf_url)
    with open(pdf_path, 'wb') as f:
        f.write(response.content)

    pages = convert_pdf_to_text(pdf_path)
    print(
    """
    Je suis Llama3, de l'équipe DREAMS TEAM, votre assistant QA pour répondre à vos questions liés aux documents 🙂 !
    Déjà pour info, le nombre de pages de vote document est: """, len(pages)
        )
    tokenizer = tiktoken.get_encoding("cl100k_base")
    chunks = split_pages_into_chunks(pages, 128, tokenizer)
    
    embedding_model = SentenceTransformer(embedding_model)

    knowledge_base = vectorize(chunks, embedding_model)
    chunks, embeddings = list(zip(*knowledge_base))
    corpus_embeddings = np.vstack(embeddings)

    llm_model = ChatOllama(model=llm_model)
    print('📑 Voici le contenu de la première page du document 😎:\n', pages[0])

    keep_looping = True
    while keep_looping:
        try:    
            question = input("Entrez votre question ✍️  |   (ou tapez 'exit' pour quitter) ✨: ")
            if question.lower() == 'exit':
                break

            response = generate_prompt(question, chunks, corpus_embeddings, embedding_model, llm_model, top_k)
            print(response)

        except KeyboardInterrupt:
           print("\nFin de la session de chat 👋.")
           keep_looping = False

if __name__ == "__main__":
    main()


"""
export CACHE_FOLDER="./cache"  #"/home/pswia/veileAI/volume_models_cache"

python chatpdf.py --pdf_url "https://diomayepresident.org/wp-content/uploads/2024/03/Programme-Diomaye-President.pdf" --embedding_model "Sahajtomar/french_semantic" --llm_model "llama3"

python chatpdf.py --pdf_url "https://arxiv.org/pdf/2302.09928" --embedding_model "Sahajtomar/french_semantic" --llm_model "llama3"




python chatpdf.py --pdf_url "https://hellofuture.orange.com/app/uploads/2024/05/2024-Orange-white-paper-on-Mobile-Network-Technology-Evolutions-Beyond-2030.pdf" --embedding_model "Sahajtomar/french_semantic" --llm_model "llama3" --top_k 5


"""