changing to nvidia nim

app.py

@@ -2,12 +2,14 @@
 import streamlit as st
 import os
 import yaml
+from langchain_nvidia_ai_endpoints import ChatNVIDIA
 from dotenv import load_dotenv
 import torch
 from src.generator import answer_with_rag
 from ragatouille import RAGPretrainedModel
 from src.data_preparation import split_documents
 from src.embeddings import init_embedding_model
+from langchain_nvidia_ai_endpoints.embeddings import NVIDIAEmbeddings
 
 from transformers import pipeline
 from langchain_community.document_loaders import PyPDFLoader
@@ -28,6 +30,7 @@ def load_config():
             return cfg
 
 cfg= load_config()
+#os.environ['NVIDIA_API_KEY']=st.secrets("NVIDIA_API_KEY")
 #load_dotenv("./src/.env")
 #HF_TOKEN=os.environ.get["HF_TOKEN"] 
 #st.write(os.environ["HF_TOKEN"] == st.secrets["HF_TOKEN"])
@@ -42,67 +45,76 @@ def main():
     st.title("Un RAG pour interroger le Collège de Pédiatrie 2024")
     user_query = st.text_input("Entrez votre question:")
 
-
-    # Initialize the retriever and LLM
+    if KNOWLEDGE_VECTOR_DATABASE not in st.session_state:
+        # Initialize the retriever and LLM
     
-    loader = PyPDFLoader(DATA_FILE_PATH)
-    #loader = PyPDFDirectoryLoader(DATA_FILE_PATH)
-    raw_document_base = loader.load()
-    MARKDOWN_SEPARATORS = [
-        "\n#{1,6} ",
-        "```\n",
-        "\n\\*\\*\\*+\n",
-        "\n---+\n",
-        "\n___+\n",
-        "\n\n",
-        "\n",
-        " ",
-        "",]
-    docs_processed = split_documents(
-        512,  # We choose a chunk size adapted to our model
-        raw_document_base,
-        tokenizer_name=EMBEDDING_MODEL_NAME,
-        separator=MARKDOWN_SEPARATORS
-    )
-    embedding_model=init_embedding_model(EMBEDDING_MODEL_NAME)
+        st.session_state.loader = PyPDFLoader(DATA_FILE_PATH)
+        #loader = PyPDFDirectoryLoader(DATA_FILE_PATH)
+        st.session_state.raw_document_base = st.session_state.loader.load()
+        st.session_state.MARKDOWN_SEPARATORS = [
+            "\n#{1,6} ",
+            "```\n",
+            "\n\\*\\*\\*+\n",
+            "\n---+\n",
+            "\n___+\n",
+            "\n\n",
+            "\n",
+            " ",
+            "",]
+        st.session_state.docs_processed = split_documents(
+            512,  # We choose a chunk size adapted to our model
+            st.session_state.raw_document_base,
+            #tokenizer_name=EMBEDDING_MODEL_NAME,
+            separator=st.session_state.MARKDOWN_SEPARATORS
+        )
+        st.session_state.embedding_model=NVIDIAEmbeddings()
+        st.session_state.KNOWLEDGE_VECTOR_DATABASE= init_vectorDB_from_doc(st.session_state.docs_processed, 
+                                                                           st.session_state.embedding_model)
     
-    if os.path.exists(VECTORDB_PATH):
-        KNOWLEDGE_VECTOR_DATABASE = FAISS.load_local(
-            VECTORDB_PATH, embedding_model, 
-            allow_dangerous_deserialization=True)
-    else:
-        KNOWLEDGE_VECTOR_DATABASE=init_vectorDB_from_doc(docs_processed, embedding_model)
-        KNOWLEDGE_VECTOR_DATABASE.save_local(VECTORDB_PATH)
+    #if os.path.exists(VECTORDB_PATH):
+      #  KNOWLEDGE_VECTOR_DATABASE = FAISS.load_local(
+      #      VECTORDB_PATH, embedding_model, 
+       #     allow_dangerous_deserialization=True)
+    #else:
+        #KNOWLEDGE_VECTOR_DATABASE=init_vectorDB_from_doc(docs_processed, embedding_model)
+       # KNOWLEDGE_VECTOR_DATABASE.save_local(VECTORDB_PATH)
 
 
     if st.button("Get Answer"):
     # Get the answer and relevant documents
-        bnb_config = BitsAndBytesConfig(
+        #bnb_config = BitsAndBytesConfig(
             #load_in_8bit=True
-            load_in_4bit=True,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-        )   
-        model = AutoModelForCausalLM.from_pretrained(READER_MODEL_NAME, quantization_config=bnb_config,
-            device_map = 'auto')
-        tokenizer = AutoTokenizer.from_pretrained(READER_MODEL_NAME)
+         #   load_in_4bit=True,
+         #   bnb_4bit_use_double_quant=True,
+         #   bnb_4bit_quant_type="nf4",
+         #   bnb_4bit_compute_dtype=torch.bfloat16,
+        #)   
+        
 
-        READER_LLM = pipeline(
-            model=model,
-            tokenizer=tokenizer,
-            task="text-generation",
-            do_sample=True,
-            temperature=0.2,
-            repetition_penalty=1.1,
-            return_full_text=False,
-            max_new_tokens=500,
-            token = os.getenv("HF_TOKEN")
+    llm = ChatNVIDIA(
+        model=READER_MODEL_NAME,
+        api_key= os.get("NVIDIA_API_KEY"), 
+        temperature=0.2,
+        top_p=0.7,
+        max_tokens=1024,
         )
-        RERANKER = RAGPretrainedModel.from_pretrained(RERANKER_MODEL_NAME)
-        num_doc_before_rerank=15
-        num_final_releveant_docs=5
-        answer, relevant_docs = answer_with_rag(query=user_query, READER_MODEL_NAME=READER_MODEL_NAME,embedding_model=embedding_model,vectorDB=KNOWLEDGE_VECTOR_DATABASE,reranker=RERANKER, llm=READER_LLM,num_doc_before_rerank=num_doc_before_rerank,num_final_relevant_docs=num_final_releveant_docs,rerank=True)
+        #tokenizer = AutoTokenizer.from_pretrained(READER_MODEL_NAME)
+
+        #READER_LLM = pipeline(
+        #    model=model,
+        #    tokenizer=tokenizer,
+        #    task="text-generation",
+        #    do_sample=True,
+       #     temperature=0.2,
+        #    repetition_penalty=1.1,
+        #    return_full_text=False,
+        #    max_new_tokens=500,
+        #    token = os.getenv("HF_TOKEN")
+       # )
+       # RERANKER = RAGPretrainedModel.from_pretrained(RERANKER_MODEL_NAME)
+       # num_doc_before_rerank=15
+       # num_final_releveant_docs=5
+       # answer, relevant_docs = answer_with_rag(query=user_query, READER_MODEL_NAME=READER_MODEL_NAME,embedding_model=embedding_model,vectorDB=KNOWLEDGE_VECTOR_DATABASE,reranker=RERANKER, llm=READER_LLM,num_doc_before_rerank=num_doc_before_rerank,num_final_relevant_docs=num_final_releveant_docs,rerank=True)
         #print(answer) 
 
     

src/generator.py

@@ -2,53 +2,47 @@
 from src.retriever import init_vectorDB_from_doc, retriever
 
 from transformers import AutoTokenizer, pipeline
+from langchain_core.prompts import ChatPromptTemplate
 from typing import List,Optional, Tuple # import the Tuple type
 from langchain.docstore.document import Document as LangchainDocument
 from langchain_community.vectorstores import FAISS
-def promt_template(query: str,READER_MODEL_NAME:str,context:str):
-    prompt_in_chat_format = [
-        {
-            "role": "system",
-            "content": """Using the information contained in the context,
-            give a comprehensive answer to the question.
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains import create_retrieval_chain
+def promt_template():
+    prompt_in_chat_format = 
+      """
+            Using the information contained in the given context, give a comprehensive answer to the question.
             Respond only to the question asked, response should be concise and relevant to the question.
-            Provide the number of the source document when relevant.If the nswer cannot be deduced from the context, do not give an answer. Please answer in french""",
-        },
-        {
-            "role": "user",
-            "content": """Context:
-            {context}
-
-            ---
-            Now here is the question you need to answer.
-
-            Question: {query}""",
-        },
-        ]
-    tokenizer = AutoTokenizer.from_pretrained(READER_MODEL_NAME)
-    RAG_PROMPT_TEMPLATE = tokenizer.apply_chat_template(
-    prompt_in_chat_format, tokenize=False, add_generation_prompt=True)
-    return RAG_PROMPT_TEMPLATE 
+            Provide the number of the source document when relevant.If the answer cannot be deduced from the context, do not give an answer. Please answer in french,
+        
+             \n\n
+            {context} """
+    prompt = ChatPromptTemplate.from_template(
+        [
+            ("system",prompt_in_chat_format),
+            ("human", "{query}")
+        ])
+    #RAG_PROMPT_TEMPLATE = tokenizer.apply_chat_template(
+    #prompt_in_chat_format, tokenize=False, add_generation_prompt=True)
+    return prompt 
 
 def answer_with_rag(
-    query: str,embedding_model, vectorDB: FAISS,READER_MODEL_NAME:str,
-    reranker,llm: pipeline,  num_doc_before_rerank: int = 5,
-    num_final_relevant_docs: int = 5,
-    rerank: bool = True
+    query: str, retriever,llm
 
 ) -> Tuple[str, List[LangchainDocument]]:
     # Build the final prompt
-    relevant_docs= retriever(query,vectorDB,reranker,num_doc_before_rerank,num_final_relevant_docs,rerank) 
-    context = "\nExtracted documents:\n"
-    context += "".join([f"Document {str(i)}:::\n" + doc for i, doc in enumerate(relevant_docs)])
+    #relevant_docs= retriever(query,vectorDB,reranker,num_doc_before_rerank,num_final_relevant_docs,rerank) 
+    #context = "\nExtracted documents:\n"
+    #context += "".join([f"Document {str(i)}:::\n" + doc for i, doc in enumerate(relevant_docs)])
     #print("=> Context:")
     #print(context)
-    RAG_PROMPT_TEMPLATE = promt_template(query,READER_MODEL_NAME,context)
-    final_prompt =RAG_PROMPT_TEMPLATE.format(query=query, context=context,READER_MODEL_NAME=READER_MODEL_NAME)
+    RAG_PROMPT_TEMPLATE = promt_template()
+    document_chain = create_stuff_documents_chain(llm, RAG_PROMPT_TEMPLATE)
+    retrieval_chain=create_retrieval_chain(retriever,document_chain)
     print("=> Final prompt:")
     #print(final_prompt)
     # Redact an answer
     print("=> Generating answer...")
-    answer = llm(final_prompt)[0]["generated_text"]
+    response=retrieval_chain.invoke({'query':query})
 
-    return answer, relevant_docs
+    return response['answer'], response["context"]