Add application file

app.py

@@ -0,0 +1,138 @@
+import gradio as gr
+import utils
+from datasets import load_dataset, concatenate_datasets
+from langchain.docstore.document import Document as LangchainDocument
+from tqdm import tqdm
+import pickle
+from ragatouille import RAGPretrainedModel
+import chunker
+import retriver
+import rag
+import nltk
+import config
+import os
+import warnings
+import sys
+import logging
+
+logging.getLogger("langchain").setLevel(logging.ERROR)
+warnings.filterwarnings("ignore")
+
+
+class AnswerSystem:
+    def __init__(self, rag_system) -> None:
+        self.rag_system = rag_system
+    
+    def answer_generate(self, question, bm_25_flag, semantic_flag, temperature):                
+        answer, relevant_docs = self.rag_system.answer(
+            question=question,
+            temperature=temperature,
+            bm_25_flag=bm_25_flag,
+            semantic_flag=semantic_flag,
+            num_retrieved_docs = 10,
+            num_docs_final = 5
+        )
+        formatted_docs = "\n\n".join([f"Document {i + 1}: {doc}" for i, doc in enumerate(relevant_docs)])
+        return answer, formatted_docs
+
+
+def run_app(rag_model):
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """
+            # RealTimeData Monthly Collection - BBC News Documentation Assistant
+
+            Welcome! This system is designed to help you explore and find insights from the RealTimeData Monthly Collection - BBC News dataset.  
+            For example:  
+            
+            - *"What position does Josko Gvardiol play, and how much did Manchester City pay for him?"*  
+
+            """
+        )
+
+        # Поля вводу
+        question_input = gr.Textbox(label="Enter your question:",
+                                    placeholder="E.g., What position does Josko Gvardiol play, and how much did Manchester City pay for him?")
+        bm25_checkbox = gr.Checkbox(label="Enable BM25-based retrieval", value=True)  # BM25 flag
+        semantic_checkbox = gr.Checkbox(label="Enable Semantic Search", value=True)  # Semantic flag
+        temperature_slider = gr.Slider(label="Response Temperature", minimum=0.1, maximum=1.0, value=0.5,
+                                       step=0.1)  # Temperature
+
+        # Кнопка пошуку
+        search_button = gr.Button("Search")
+
+        # Поля виводу
+        answer_output = gr.Textbox(label="Answer", interactive=False, lines=5)
+        docs_output = gr.Textbox(label="Relevant Documents", interactive=False, lines=10)
+
+        # Логіка пошуку
+        system = AnswerSystem(rag_model)
+
+        search_button.click(
+            system.answer_generate,
+            inputs=[question_input, bm25_checkbox, semantic_checkbox, temperature_slider],  # Всі параметри
+            outputs=[answer_output, docs_output]
+        )
+
+    # Запуск додатку
+    demo.launch(debug=True, share=True)
+
+
+def get_rag_data():
+    nltk.download('punkt')
+    nltk.download('punkt_tab')
+
+    if os.path.exists(config.DOCUMENTS_PATH):
+        print(f"Loading preprocessed documents from {config.DOCUMENTS_PATH}")
+        with open(config.DOCUMENTS_PATH, "rb") as file:
+            docs_processed = pickle.load(file)
+    else:
+        print("Processing documents...")
+        datasets_list = [
+            utils.align_features(load_dataset("RealTimeData/bbc_news_alltime", config)["train"])
+            for config in tqdm(config.AVAILABLE_DATASET_CONFIGS)
+        ]
+
+        ds = concatenate_datasets(datasets_list)
+
+        RAW_KNOWLEDGE_BASE = [
+            LangchainDocument(
+                page_content=doc["content"],
+                metadata={
+                    "title": doc["title"],
+                    "published_date": doc["published_date"],
+                    "authors": doc["authors"],
+                    "section": doc["section"],
+                    "description": doc["description"],
+                    "link": doc["link"]
+                }
+            )
+            for doc in tqdm(ds)
+        ]
+
+        docs_processed = chunker.split_documents(512, RAW_KNOWLEDGE_BASE)
+
+        print(f"Saving preprocessed documents to {config.DOCUMENTS_PATH}")
+        with open(config.DOCUMENTS_PATH, "wb") as file:
+            pickle.dump(docs_processed, file)
+
+    return docs_processed
+
+
+if __name__ == '__main__':
+    docs_processed = get_rag_data()
+
+    bm25 = retriver.create_bm25(docs_processed)
+
+    KNOWLEDGE_VECTOR_DATABASE = retriver.create_vector_db(docs_processed)
+
+    RERANKER = RAGPretrainedModel.from_pretrained(config.CROSS_ENCODER_MODEL)
+
+    rag_generator = rag.RAGAnswerGenerator(
+        docs=docs_processed,
+        bm25=bm25,
+        knowledge_index=KNOWLEDGE_VECTOR_DATABASE,
+        reranker=RERANKER
+    )
+
+    run_app(rag_generator)

chunker.py

@@ -0,0 +1,45 @@
+import config
+from langchain.docstore.document import Document as LangchainDocument
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from transformers import AutoTokenizer
+from tqdm.notebook import tqdm
+from typing import List
+
+
+def split_documents(chunk_size: int, knowledge_base: List[LangchainDocument]) -> List[LangchainDocument]:
+    """
+    Split documents into chunks of maximum size `chunk_size` tokens and return a list of documents.
+    """
+    MARKDOWN_SEPARATORS = [
+        "\n#{1,6} ",
+        "```\n",
+        "\n\\*\\*\\*+\n",
+        "\n---+\n",
+        "\n___+\n",
+        "\n\n",
+        "\n",
+        " ",
+        "",
+    ]
+
+    text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
+        AutoTokenizer.from_pretrained(config.EMBEDDING_MODEL_NAME),
+        chunk_size=chunk_size,
+        chunk_overlap=int(chunk_size / 10),
+        add_start_index=True,
+        strip_whitespace=True,
+        separators=MARKDOWN_SEPARATORS,
+    )
+
+    docs_processed = []
+    for doc in tqdm(knowledge_base):
+        docs_processed += text_splitter.split_documents([doc])
+
+    unique_texts = {}
+    docs_processed_unique = []
+    for doc in docs_processed:
+        if doc.page_content not in unique_texts:
+            unique_texts[doc.page_content] = True
+            docs_processed_unique.append(doc)
+
+    return docs_processed_unique # , docs_processed

config.py

@@ -0,0 +1,28 @@
+import os
+
+
+RAG_PROMPT = """
+              You are an advanced Retrieval-Augmented Generation (RAG) Assistant.
+
+              Your task is to answer user questions based only on the provided documents. Use the context from the documents to generate a response.
+
+              **Guidelines:**
+              1. **Always cite sources**: When information is derived from a document, reference it by citing the chunk number in square brackets, e.g., [Chunk 1], where relevant information is used.
+              2. If the answer cannot be determined from the provided documents, state: "The answer cannot be determined from the provided documents."
+              3. After each answer, provide a numbered list of the retrieved chunks.
+
+              Please follow these instructions to generate accurate and well-cited answers based on the documents.
+              """
+LLM_ONLY_PROMPT = """You are an Assistant. If no documents are retrieved, answer the question based on general knowledge."""
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+os.environ['GROQ_API_KEY'] = "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"
+DB_PATH = "vector_database.faiss"
+BM25_PATH = "bm25_index.pkl"
+DOCUMENTS_PATH = "processed_documents.pkl"
+EMBEDDING_MODEL_NAME = "thenlper/gte-small"
+CROSS_ENCODER_MODEL = "colbert-ir/colbertv2.0"
+
+AVAILABLE_DATASET_CONFIGS = [
+    '2024-11'
+]

rag.py

@@ -0,0 +1,100 @@
+from typing import Optional, List, Tuple
+from langchain.docstore.document import Document as LangchainDocument
+from rank_bm25 import BM25Okapi
+from langchain_community.vectorstores import FAISS
+from ragatouille import RAGPretrainedModel
+from litellm import completion
+import os
+import retriver
+import config
+
+
+class RAGAnswerGenerator:
+    def __init__(self, docs: List[LangchainDocument], bm25: BM25Okapi, knowledge_index: FAISS, reranker: Optional[RAGPretrainedModel] = None):
+        self.bm25 = bm25
+        self.knowledge_index = knowledge_index
+        self.docs = docs
+        self.reranker = reranker
+        self.llm_key = os.environ['GROQ_API_KEY']
+
+    def retrieve_documents(
+        self,
+        question: str,
+        num_retrieved_docs: int,
+        bm_25_flag: bool,
+        semantic_flag: bool
+    ) -> List[str]:
+        print("=> Retrieving documents...")
+        relevant_docs = []
+
+        if bm_25_flag or semantic_flag:
+            result = retriver.search(
+                self.docs,
+                self.bm25,
+                self.knowledge_index,
+                question,
+                use_bm25=bm_25_flag,
+                use_semantic_search=semantic_flag,
+                top_k=num_retrieved_docs
+            )
+            if bm_25_flag and semantic_flag:
+                relevant_docs = [doc.page_content for doc in result]
+                return relevant_docs
+            elif bm_25_flag:
+                relevant_docs = result
+                return relevant_docs
+            elif semantic_flag:
+                relevant_docs = [doc.page_content for doc in result]
+                return relevant_docs
+                
+
+    def rerank_documents(self, question: str, documents: List[str], num_docs_final: int) -> List[str]:
+        if self.reranker and documents:
+            print("=> Reranking documents...")
+            reranked_docs = self.reranker.rerank(question, documents, k=num_docs_final)
+            return [doc["content"] for doc in reranked_docs]
+        return documents[:num_docs_final]
+
+    def format_context(self, documents: List[str]) -> str:
+        if not documents:
+            return "No retrieved documents available."
+        return "\n".join([f"[{i + 1}] {doc}" for i, doc in enumerate(documents)])
+
+    def generate_answer(
+        self,
+        question: str,
+        context: str,
+        temperature: float,
+    ) -> str:
+        print("=> Generating answer...")
+        if context.strip() == "No retrieved documents available.":
+            response = completion(
+                model="groq/llama3-8b-8192",
+                messages=[
+                    {"role": "system", "content": config.LLM_ONLY_PROMPT},
+                    {"role": "user", "content": f"Question: {question}"}
+                ],
+                api_key=self.llm_key,
+                temperature=temperature
+            )
+        else:
+            response = completion(
+                model="groq/llama3-8b-8192",
+                messages=[
+                    {"role": "system", "content": config.RAG_PROMPT},
+                    {"role": "user", "content": f""" Context: {context} Question: {question} """}
+                ],
+                api_key=self.llm_key,
+                temperature=temperature
+            )
+        return response.get("choices", [{}])[0].get("message", {}).get("content", "No response content found")
+
+    def answer(self, question: str, temperature: float, num_retrieved_docs: int = 30, num_docs_final: int = 5, bm_25_flag=True, semantic_flag=True) -> Tuple[str, List[str]]:
+        relevant_docs = self.retrieve_documents(question, num_retrieved_docs, bm_25_flag, semantic_flag)
+        print(len(relevant_docs))
+        relevant_docs = self.rerank_documents(question, relevant_docs, num_docs_final)
+        print(len(relevant_docs))
+        context = self.format_context(relevant_docs)
+        answer = self.generate_answer(question, context, temperature)
+        document_list = [f"[{i + 1}] {doc}" for i, doc in enumerate(relevant_docs)] if relevant_docs else []
+        return answer, document_list

retriver.py

@@ -0,0 +1,96 @@
+import config
+import utils
+from nltk.tokenize import word_tokenize
+from typing import List
+import nltk
+import torch
+import pickle
+from langchain.docstore.document import Document as LangchainDocument
+from rank_bm25 import BM25Okapi
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores.utils import DistanceStrategy
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
+import os
+
+
+def create_vector_db(docs: List[LangchainDocument]):
+    db_path: str = config.DB_PATH
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    embedding_model = HuggingFaceEmbeddings(
+        model_name=config.EMBEDDING_MODEL_NAME,
+        multi_process=True,
+        model_kwargs={"device": device},
+        encode_kwargs={"normalize_embeddings": True},
+    )
+    
+    if os.path.exists(db_path):
+        print(f"Завантаження векторної бази даних з {db_path}")
+        knowledge_vector_database = FAISS.load_local(
+            db_path, 
+            embedding_model,
+            allow_dangerous_deserialization=True
+        )
+        return knowledge_vector_database
+    elif docs is not None:
+        print("Створення нової векторної бази даних")
+        knowledge_vector_database = FAISS.from_documents(
+            docs, embedding_model, distance_strategy=DistanceStrategy.COSINE
+        )
+        knowledge_vector_database.save_local(db_path)
+        print(f"Векторна база даних збережена в {db_path}")
+        return knowledge_vector_database
+    else:
+      raise ValueError(
+            """Documents are missing! 
+            Please load the documents and set get_data=True in app.py."""
+        )
+    
+
+
+def create_bm25(docs: List[LangchainDocument]):
+    bm25_path: str = config.BM25_PATH
+    if os.path.exists(bm25_path):
+        print(f"Завантаження BM25 індексу з {bm25_path}")
+        with open(bm25_path, "rb") as file:
+            bm25 = pickle.load(file)
+        return bm25
+    elif docs is not None:
+        print("Створення нового BM25 індексу")
+        tokenized_docs = [word_tokenize(doc.page_content.lower()) for doc in docs]
+        bm25 = BM25Okapi(tokenized_docs)
+        with open(bm25_path, "wb") as file:
+            pickle.dump(bm25, file)
+        print(f"BM25 індекс збережено в {bm25_path}")
+        return bm25
+    else:
+      raise ValueError(
+            """Documents are missing! 
+            Please load the documents and set get_data=True in app.py."""
+        )
+
+
+def search(docs_processed, bm_25: BM25Okapi, vector_db: FAISS, query, top_k, use_bm25=True, use_semantic_search=True):
+    if use_bm25 and use_semantic_search:
+        bm25_retriever = BM25Retriever.from_documents(docs_processed)
+        bm25_retriever.k = top_k
+        faiss_retriever = vector_db.as_retriever(search_kwargs={"k": top_k})
+        ensemble_retriever = EnsembleRetriever(
+            retrievers=[bm25_retriever, faiss_retriever],
+            weights=[0.5, 0.5]
+        )      
+        result = ensemble_retriever.invoke(query)
+        return result
+
+    elif use_bm25:
+        tokenized_query = word_tokenize(query.lower())
+        result = bm_25.get_top_n(tokenized_query, [doc.page_content for doc in docs_processed], n=top_k)
+
+    elif use_semantic_search:
+        result = vector_db.similarity_search(query, k=top_k)
+    else:
+       result = []
+    return result
+    

utils.py

@@ -0,0 +1,12 @@
+import config
+from transformers import AutoTokenizer
+from typing import List
+
+
+def align_features(dataset):
+    def fix_authors(example):
+        if not isinstance(example["authors"], list):
+             return {"authors": [example["authors"]] if example["authors"] else []}
+        return example
+        
+    return dataset.map(fix_authors)