fixing main

app.py

@@ -10,7 +10,7 @@ from src.embeddings import init_embedding_model
 
 from transformers import pipeline
 from langchain_community.document_loaders import PyPDFLoader
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from src.retriever import init_vectorDB_from_doc, retriever
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from langchain_community.vectorstores import FAISS
@@ -36,77 +36,82 @@ READER_MODEL_NAME=cfg['READER_MODEL_NAME']
 RERANKER_MODEL_NAME=cfg['RERANKER_MODEL_NAME']
 VECTORDB_PATH=cfg['VECTORDB_PATH']
 
-st.title("Une application RAG pour interroger le Collège de Pédiatrie 2024")
 
-user_query = st.text_input("Entrez votre question:")
+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
     
-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)
+    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)
     
-if os.path.exists(VECTORDB_PATH):
-    new_vector_store = 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):
+        new_vector_store = 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"):
+    if st.button("Get Answer"):
     # Get the answer and relevant documents
-    bnb_config = BitsAndBytesConfig(
-        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)
-    tokenizer = AutoTokenizer.from_pretrained(READER_MODEL_NAME)
+        bnb_config = BitsAndBytesConfig(
+            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)
+        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) 
+        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) 
 
     
-    # Display the answer
-    st.write("### Answer:")
-    st.write(answer)
+        # Display the answer
+        st.write("### Answer:")
+        st.write(answer)
     
-    # Display the relevant documents
-    st.write("### Relevant Documents:")
-    for i, doc in enumerate(relevant_docs):
-        st.write(f"Document {i}:\n{doc.text}")
+        # Display the relevant documents
+        st.write("### Relevant Documents:")
+        for i, doc in enumerate(relevant_docs):
+            st.write(f"Document {i}:\n{doc.text}")
+
+
+if __name__ == "__main__":
+    main()

src/embeddings.py

@@ -1,5 +1,5 @@
 # Databricks notebook source
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores.utils import DistanceStrategy