update

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Tran Anh Quoc
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,66 @@
----
-title: GPT2 Vietnamese
-emoji: 🚀
-colorFrom: gray
-colorTo: green
-sdk: gradio
-sdk_version: 4.21.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# RAG-Based Chatbot System
+
+## Overview
+This project aims to build a Retrieval-Augmented Generation (RAG)-based chatbot system. The chatbot utilizes context-aware chunking for efficient document processing and leverages open-source models for embeddings and language generation.
+
+## Features
+- **Context-Aware Chunking**: Implements an optimal, manual chunking strategy that uses special chunk markers to separate chunks within documents. This allows for easy extraction of chunks using Python's split() method. Make sure the documents users upload have been chunked using a special chunk marker separator for effective processing.
+  
+- **Data Processing Tool**: - **Data Processing Tool**: Converts tables in documents into HTML tables to handle challenges such as long tables and merged cells. The tool requires users to bold headers in tables for clarity and automatically identifies table types, including those with 1 header or more than 2 headers.
+
+- **Open-Source Models**: Uses open-source models instead of proprietary ones like those from OpenAI, providing a cost-effective and flexible solution.
+  - **Embedding Model**: Utilizes `intfloat/multilingual-e5-small`, which is highly efficient and particularly effective for Vietnamese text.
+  - **Language Model**: Uses `Viet-Mistral/Vistral-7B-Chat`, a language model based on Mistral, with continued pretraining on Vietnamese for better generation performance.
+
+## Installation
+1. Clone the repository:
+```sh
+git clone https://github.com/quoctata2911/RAG-based-ChatBot-System.git
+```
+
+2. Navigate to the project directory:
+```sh
+cd RAG-Based-Chatbot-System
+```
+
+3. Install the required dependencies:
+```sh
+pip install -r requirements.txt
+```
+
+## Usage
+Upload your Word .docx documents into the data folder. Ensure that each document has been chunked using a special chunk marker separator as specified in the config.yaml file.
+
+1. Configure the chunk marker:
+- Open the `config.yaml` file located in the project directory.
+- Locate the parameter defining the chunk marker and adjust it as needed for your document segmentation requirements.
+
+2. Prepare the data:
+```sh
+python prepare_data.py
+```
+3. Run the chatbot:
+```sh
+python chat.py
+```
+
+## Project Structure
+- **prepare_data.py**: Script to preprocess and chunk documents, converting tables into HTML and segmenting them with chunk markers.
+- **chat.py**: Main script to run the chatbot system.
+
+## Models
+- **Embedding Model**: We use the `intfloat/multilingual-e5-small` model for generating embeddings. This model is particularly effective for Vietnamese text, outperforming other models in our benchmarks.
+
+- **Language Model**: The language model used is Vistral, a variant of the Mistral model that has been further pre-trained on Vietnamese text for improved performance in language generation tasks.
+
+## Benchmarking and Performance
+Through extensive benchmarking, the `intfloat/multilingual-e5-small` model has proven to be the best choice for Vietnamese embeddings, offering a balance of efficiency and performance. The Vistral model enhances language generation capabilities, ensuring the chatbot responds accurately and naturally in Vietnamese.
+
+## Contributions
+We welcome contributions to improve the RAG-ChatBot. Please fork the repository and create a pull request with your changes. For major changes, please open an issue first to discuss what you would like to change.
+
+## License
+This project is licensed under the MIT License. See the LICENSE file for more details.
+
+## Contact
+For any questions or suggestions, please contact me at quoctrananh2911@gmail.com

app.py

@@ -1,50 +1,192 @@
-import gradio as gr
+import yaml
 import torch
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import logging
+import argparse
+import warnings
+import pandas as pd
+from tqdm.auto import tqdm
+from jsonargparse import CLI
+from types import SimpleNamespace
+from llama_index.core.schema import TextNode
+from langchain_huggingface import HuggingFaceEmbeddings
+from llama_index.core import Prompt, Settings, VectorStoreIndex
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer
+
+import gradio as gr
+
+
+
+def load_config(config_path='config.yaml'):
+    print('-> Loading config file ...')
+    cfg = yaml.safe_load(
+        open(config_path).read()
+    )
+
+    for k,v in cfg.items():
+        if type(v) == dict:
+            cfg[k] = SimpleNamespace(**v)
+    cfg = SimpleNamespace(**cfg)
+    return cfg
+
+def get_prompt_template():
+    template = (
+    "Bạn là trợ lý ảo hữu ích và thông minh được huấn luyên được để trả lời các câu hỏi từ người dùng giữa trên các thông tin ngữ cảnh liên quan được cung cấp\n"
+    "Thông tin ngữ cảnh:\n"
+    "---------------------\n"
+    "{context_str}"
+    "\n---------------------\n"
+    "Dựa trên những thông tin ngữ cảnh bên trên, hãy trả lời câu hỏi sau: {query_str}\n"
+    )
+    qa_template = Prompt(template)
+    return qa_template
+
+def reset_settings(cfg):
+    embed_model =HuggingFaceEmbeddings(
+        model_name=cfg.architecture.embedding_model
+    )
+    Settings.embed_model = embed_model
+    Settings.llm = None  
 
+def get_retriever(cfg, prompt_template):
+    chunks = pd.read_pickle('processed_chunks.pickle')['chunk'].values.tolist()
+    nodes = [TextNode(text=chunk) for chunk in chunks]
+    index = VectorStoreIndex(nodes=nodes)
+    retriever = index.as_query_engine(
+        similarity_top_k=cfg.retrieve.top_k,
+        text_qa_template=prompt_template
+    )
+    return retriever
 
-tokenizer = GPT2Tokenizer.from_pretrained('NlpHUST/gpt2-vietnamese')
-model = GPT2LMHeadModel.from_pretrained('NlpHUST/gpt2-vietnamese')
-# max_length = 100
+def load_tokenizer(cfg):
+    tokenizer =  AutoTokenizer.from_pretrained(
+        cfg.architecture.llm_model,
+        token=cfg.architecture.hf_token
+    )
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer
+
+def get_llm(cfg):
+    if cfg.architecture.llm_quantized:
+        bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.float16
+            )
+    else:
+        bnb_config = None
+            
+
+    llm = AutoModelForCausalLM.from_pretrained(
+        cfg.architecture.llm_model,
+        torch_dtype=torch.bfloat16,
+        device_map=cfg.environment.device,
+        token=cfg.architecture.hf_token,
+        low_cpu_mem_usage=True,
+        quantization_config=bnb_config,
+    )
+
+    return llm.eval()
 
 
 def run(text, intensity):
-    res="Tham khảo NlpHUST model  \n  \n  \n"
+    prompt = retriever.query(text).response
+    prompt = tokenizer.bos_token + '[INST] ' + prompt + ' [/INST]'
+    streamer = TextStreamer(tokenizer, skip_prompt=True)
+    input_ids = tokenizer([prompt], return_tensors='pt').to(cfg.environment.device)
+
+    _ = language_model.generate(
+            **input_ids,
+            streamer=streamer,
+            pad_token_id=tokenizer.pad_token_id,
+            max_new_tokens=cfg.generation.max_new_tokens,
+            do_sample=cfg.generation.do_sample,
+            temperature=cfg.generation.temperature
+        )
+
+    # print(20*'---')
+
+
+
+    res="Chatbot Data Mining 2024   \n  \n  \n"
     max_length=intensity
-        
-    input_ids = tokenizer.encode(text, return_tensors='pt')
-    sample_outputs = model.generate(input_ids,pad_token_id=tokenizer.eos_token_id,
-                                   do_sample=True,
-                                   max_length=max_length,
-                                   min_length=5,
-                                   top_k=40,
-                                   num_beams=5,
-                                   early_stopping=True,
-                                   no_repeat_ngram_size=2,
-                                   num_return_sequences=2)
-    
-    for i, sample_output in enumerate(sample_outputs):
-        res +="Mẫu số {}\n \n{}".format(i+1, tokenizer.decode(sample_output.tolist()))
-        res +='\n \n \n \n'
-    return res
-
-# demo = gr.Interface(
-#     fn=run,
-#     inputs=["text", "slider"],
-#     outputs=["text"],
-# )
-
-demo = gr.Interface(fn=run,
+
+    return _
+
+
+def vistral_chat(cfg, retriever, tokenizer, language_model):
+    demo = gr.Interface(fn=run,
                     inputs=[gr.Textbox(label="Nhập vào nội dung input",value="Con đường xưa em đi"),gr.Slider(label="Độ dài output muốn tạo ra", value=20, minimum=10, maximum=100, step=2)],
                     outputs=gr.Textbox(label="Output"),  # <-- Number of output components: 1
                     )
 
-demo.launch()
+    demo.launch()
 
 
 
+    # while True:
+    #     user_query = input('👨‍🦰 ')
+    #     prompt = retriever.query(user_query).response
+    #     prompt = tokenizer.bos_token + '[INST] ' + prompt + ' [/INST]'
+    #     streamer = TextStreamer(tokenizer, skip_prompt=True)
+    #     input_ids = tokenizer([prompt], return_tensors='pt').to(cfg.environment.device)
 
+    #     _ = language_model.generate(
+    #         **input_ids,
+    #         streamer=streamer,
+    #         pad_token_id=tokenizer.pad_token_id,
+    #         max_new_tokens=cfg.generation.max_new_tokens,
+    #         do_sample=cfg.generation.do_sample,
+    #         temperature=cfg.generation.temperature
+    #     )
 
+    #     print(20*'---')
 
 
+def main(config_path):
+    # Configure logging
+    logging.basicConfig(level=logging.INFO, 
+                        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    logger = logging.getLogger(__name__)
+    
+    try:
+        # Log the start of the process
+        logger.info("Starting the process with config file: %s", config_path)
+        
+        # Load configuration from the file
+        config = load_config(config_path)
+        
+        # Load necessary components
+        prompt_template = get_prompt_template()
+        
+        # Replace OpenAI embed model and llm with custom ones
+        reset_settings(config)
+        
+        # Get retriever
+        retriever = get_retriever(config, prompt_template)
+        
+        # Load tokenizer and language model
+        tokenizer = load_tokenizer(config)
+        language_model = get_llm(config)
+        
+        # Start the command line interface
+        vistral_chat(config, retriever, tokenizer, language_model)
+        
+        # Log successful completion
+        logger.info("Process completed successfully.")
+        
+    except FileNotFoundError as e:
+        logger.error("Configuration file not found: %s", e)
+    except Exception as e:
+        logger.exception("An error occurred: %s", e)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process some configurations.')
+    parser.add_argument('--config', type=str, default='config.yaml', help='Path to the configuration file')
+    args = parser.parse_args()
+    main(args.config)
+
+    
 

app1.py

@@ -0,0 +1,50 @@
+import gradio as gr
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+
+
+tokenizer = GPT2Tokenizer.from_pretrained('NlpHUST/gpt2-vietnamese')
+model = GPT2LMHeadModel.from_pretrained('NlpHUST/gpt2-vietnamese')
+# max_length = 100
+
+
+def run(text, intensity):
+    res="Tham khảo NlpHUST model  \n  \n  \n"
+    max_length=intensity
+        
+    input_ids = tokenizer.encode(text, return_tensors='pt')
+    sample_outputs = model.generate(input_ids,pad_token_id=tokenizer.eos_token_id,
+                                   do_sample=True,
+                                   max_length=max_length,
+                                   min_length=5,
+                                   top_k=40,
+                                   num_beams=5,
+                                   early_stopping=True,
+                                   no_repeat_ngram_size=2,
+                                   num_return_sequences=2)
+    
+    for i, sample_output in enumerate(sample_outputs):
+        res +="Mẫu số {}\n \n{}".format(i+1, tokenizer.decode(sample_output.tolist()))
+        res +='\n \n \n \n'
+    return res
+
+# demo = gr.Interface(
+#     fn=run,
+#     inputs=["text", "slider"],
+#     outputs=["text"],
+# )
+
+demo = gr.Interface(fn=run,
+                    inputs=[gr.Textbox(label="Nhập vào nội dung input",value="Con đường xưa em đi"),gr.Slider(label="Độ dài output muốn tạo ra", value=20, minimum=10, maximum=100, step=2)],
+                    outputs=gr.Textbox(label="Output"),  # <-- Number of output components: 1
+                    )
+
+demo.launch()
+
+
+
+
+
+
+
+

config.yaml

@@ -0,0 +1,21 @@
+architecture:
+    llm_model: Viet-Mistral/Vistral-7B-Chat
+    embedding_model: intfloat/multilingual-e5-small
+    hf_token: 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
+    llm_quantized: False
+dataset:
+    chunk_marker: BK_CHUNK
+    required_exts: .docx
+    data_dir: ./data
+    processed_data_dir: ./processed_data
+    signal_type: html_table
+    keep_bold: True
+retrieve:
+    top_k: 2
+generation:
+    max_new_tokens: 2048
+    temperature: 0.6
+    do_sample: True
+    top_p: 0.9
+environment:
+    device: 'cuda'

prepare_data.py

@@ -0,0 +1,100 @@
+import os
+import shutil
+import yaml
+import logging
+import pandas as pd
+from pathlib import Path
+from jsonargparse import CLI
+from docx.api import Document
+from types import SimpleNamespace
+from llama_index.core import SimpleDirectoryReader
+from utils.process_tables import extract_and_replace_docx_tables
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("script.log"),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+def load_config(file_path='config.yaml'):
+    logger.info('Loading config file ...')
+    try:
+        with open(file_path, 'r') as file:
+            cfg = yaml.safe_load(file)
+        for k, v in cfg.items():
+            if isinstance(v, dict):
+                cfg[k] = SimpleNamespace(**v)
+        logger.info('Config file loaded successfully.')
+        return SimpleNamespace(**cfg)
+    except Exception as e:
+        logger.error(f'Error loading config file: {e}')
+        raise
+
+cfg = load_config()
+
+def process_docx_files(data_dir=Path(cfg.dataset.data_dir), 
+                       processed_data_dir=Path(cfg.dataset.processed_data_dir),
+                       chunk_marker=cfg.dataset.chunk_marker):
+    try:
+        if not os.path.exists(processed_data_dir):
+            shutil.rmtree(processed_data_dir)
+
+        docx_files = [file for file in os.listdir(data_dir) if file.endswith('.docx')]
+        logger.info(f'Found {len(docx_files)} DOCX files to process.')
+
+        for fname in docx_files:
+            document, html_chunked_tables = extract_and_replace_docx_tables(
+                docx_file=data_dir / fname, 
+                chunk_marker=chunk_marker
+            )
+            document.save(processed_data_dir / f'processed_{fname}')
+            logger.info(f'Processed and saved {fname}')
+    except Exception as e:
+        logger.error(f'Error processing DOCX files: {e}')
+        raise
+
+def load_processed_data(processed_data_dir=Path(cfg.dataset.processed_data_dir)):
+    try:
+        documents = SimpleDirectoryReader(
+            input_dir=processed_data_dir,
+            required_exts=[cfg.dataset.required_exts],
+        ).load_data()
+        logger.info('Processed data loaded successfully.')
+        return documents
+    except Exception as e:
+        logger.error(f'Error loading processed data: {e}')
+        raise
+
+def get_chunks(documents, chunk_marker=cfg.dataset.chunk_marker):
+    try:
+        chunks = [chunk.strip() for doc in documents for chunk in doc.text.split(chunk_marker) if chunk.strip()]
+        logger.info(f'Extracted {len(chunks)} chunks from documents.')
+        return chunks
+    except Exception as e:
+        logger.error(f'Error extracting chunks: {e}')
+        raise
+
+def main():
+    logger.info('Starting document processing ...')
+    try:
+        process_docx_files()
+
+        documents = load_processed_data()
+        chunks = get_chunks(documents)
+        num_chunks = len(chunks)
+        logger.info(f'Total number of chunks: {num_chunks}')
+
+        df_chunks = pd.DataFrame({'chunk': chunks})
+        df_chunks.to_pickle('processed_chunks.pickle')
+        logger.info('All chunks saved to processed_chunks.pickle')
+    except Exception as e:
+        logger.error(f'Error in main processing: {e}')
+        raise
+
+if __name__ == '__main__':
+    main()

processed_chunks.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c48413343bfeff0b6cff53fcb2328358c2b8d0e1006f89b4ebd44e679ec96bb2
+size 33042

requirements.txt

@@ -1,2 +1,12 @@
-torch
-transformers
+llama-index==0.10.44
+mammoth==1.7.1
+python-docx==1.1.2
+docx2txt==0.8
+langchain-community==0.2.4
+sentence-transformers==3.0.1
+accelerate==0.31.0
+bitsandbytes==0.43.1
+langchain-huggingface==0.0.3
+jsonargparse==4.29.0
+llama-index-embeddings-langchain
+gradio

utils/process_tables.py

@@ -0,0 +1,141 @@
+import time
+import mammoth
+from bs4 import BeautifulSoup
+from docx.api import Document
+import re
+
+def extract_and_replace_docx_tables(docx_file, chunk_marker):
+    start_time = time.time()  # Record start time
+
+    document = Document(docx_file)
+    docx_tables = document.tables
+    total_tables = len(document.tables)
+
+    with open(docx_file, "rb") as docx_file:
+        result = mammoth.convert_to_html(docx_file)
+        html = result.value 
+
+    tables = extract_html_tables(html)
+    
+    html_chunked_tables = get_html_table_chunks(tables, chunk_marker=chunk_marker)
+    
+    
+    html_tables = []
+    for table in html_chunked_tables:
+        temp_document = Document()
+        html_table = temp_document.add_paragraph(table)._element
+        html_table.alignment = 0
+        html_tables.append(html_table)
+
+
+    track = 0    
+    while len(document.tables) > 0:
+        track += 1
+        try:
+            html_table = html_tables[0]
+            document.element.body.replace(document.tables[0]._element, html_table)
+            html_tables.remove(html_table)
+            end_time = time.time()  # Record end time
+            # print(f'{track} of {total_tables} | Success | Time: {end_time - start_time:.2f} seconds')
+
+        except Exception as e:
+            print(f'{track} of {total_tables} | Fail: {e}')
+            if track >= 200:
+                break
+            
+    return document, html_chunked_tables
+
+def extract_html_tables(html):
+    soup = BeautifulSoup(html, 'html.parser')
+    tables = soup.find_all('table')
+    return tables
+
+def get_html_table_chunks(tables, chunk_marker):
+    
+    html_chunk_marker = '<strong>' + chunk_marker + '</strong>'
+
+    html_table_chunks = []
+    
+    for table_soup in tables:
+
+        html_table_string = str(table_soup)
+        html_table_string = html_table_string.replace('<table>', '<table>\n')
+        html_table_string = html_table_string.replace('<tr>', '\n<tr>\n')
+        html_table_string = html_table_string.replace('</tr>', '\n</tr>')
+        html_table_string = html_table_string.replace('<thead>', '<thead>')
+        html_table_string = html_table_string.replace('</thead>', '\n</thead>\n')
+        html_table_string = html_table_string.replace('<tbody>', '<tbody>')
+        html_table_string = html_table_string.replace('</tbody>', '\n</tbody>\n')
+
+        with open('table_html.txt', mode='w', encoding='utf8') as f:
+            f.write(html_table_string)
+
+        with open('table_html.txt', mode='r', encoding='utf8') as f:
+            lines = f.readlines()
+
+        start_table = lines[0].strip()
+        end_table = lines[-1].strip()
+        
+        # Get start and end tags for tbody
+        start_tbody = '<tbody>' if '<tbody>' in html_table_string else ''
+        end_tbody = '</tbody>' if '</tbody>' in html_table_string else ''
+
+        # Extract and clean headers if present
+        headers = str(table_soup.find('thead')) if 'thead' in html_table_string else ''
+        headers = re.sub(r'>\n\s*<', '><', headers)
+
+        processed_lines = []
+        for line in lines:
+            if chunk_marker in line:
+                start_index = line.find(html_chunk_marker)
+                chunk_start = start_index - len('<p>')
+                chunk_end = start_index + len(html_chunk_marker) + len('</p>')
+
+                chunk_html = line[chunk_start:chunk_end]
+                if chunk_html.startswith('<p>') & chunk_html.endswith('</p>'):
+                    line = line.replace('<p>', '')
+                    line = line.replace('</p>', '')
+                else:
+                    pass
+
+                line = line.replace(html_chunk_marker, '')
+                line = line.replace(' </td>', '</td>').strip()
+                line += chunk_marker
+
+            processed_lines.append(line)
+
+        processed_lines = [line.strip() for line in processed_lines]
+        html_table = ''.join(processed_lines)
+        html_chunks = html_table.split(chunk_marker)
+
+        proccessed_html_chunks = []
+        for index, chunk in enumerate(html_chunks):
+            if index == 0:
+                chunk += (end_tbody + end_table)
+                first_chunk = chunk.replace(end_table, '')
+                start = first_chunk.find('<tr>')
+                end = first_chunk.find('</tr>') + len('</tr>')
+                headers = first_chunk[start:end]
+
+                    
+            elif chunk == html_chunks[-1]:
+                chunk = start_table + headers + start_tbody + chunk 
+
+            else:
+                chunk = start_table + headers + start_tbody + chunk + end_tbody + end_table
+
+            proccessed_html_chunks.append(chunk)
+
+        chunks_to_html = ''
+        for html_chunk in proccessed_html_chunks:
+                chunks_to_html += wrap_signal(html_chunk, signal_type='html')
+                if html_chunk != proccessed_html_chunks[-1]:
+                    chunks_to_html += f'\n\n{chunk_marker}\n\n'
+        
+        html_table_chunks.append(chunks_to_html)
+        
+    return html_table_chunks
+
+def wrap_signal(data, signal_type):
+    data = f"```{signal_type}\n{data}\n```"
+    return data