Update app.py

app.py

@@ -1,4 +1,245 @@
 # @title web interface demo
+# import random
+# import gradio as gr
+# import time
+# import numpy as np
+# import pandas as pd
+# import torch
+# import faiss
+# from sklearn.preprocessing import normalize
+# from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+# from sentence_transformers import SentenceTransformer, util
+# from pythainlp import Tokenizer
+# import pickle
+# import re
+# from pythainlp.tokenize import sent_tokenize
+
+# DEFAULT_MODEL = 'wangchanberta-hyp'
+# DEFAULT_SENTENCE_EMBEDDING_MODEL = 'intfloat/multilingual-e5-base'
+
+# MODEL_DICT = {
+#     'wangchanberta': 'Chananchida/wangchanberta-xet_ref-params',
+#     'wangchanberta-hyp': 'Chananchida/wangchanberta-xet_hyp-params',
+# }
+
+# EMBEDDINGS_PATH = 'data/embeddings.pkl'
+# DATA_PATH='data/dataset.xlsx'
+
+
+# def load_data(path=DATA_PATH):
+#     df = pd.read_excel(path, sheet_name='Default')
+#     df['Context'] = pd.read_excel(path, sheet_name='mdeberta')['Context']
+#     print(len(df))
+#     print('Load data done')
+#     return df
+
+
+# def load_model(model_name=DEFAULT_MODEL):
+#     model = AutoModelForQuestionAnswering.from_pretrained(MODEL_DICT[model_name])
+#     tokenizer = AutoTokenizer.from_pretrained(MODEL_DICT[model_name])
+#     print('Load model done')
+#     return model, tokenizer
+
+# def load_embedding_model(model_name=DEFAULT_SENTENCE_EMBEDDING_MODEL):
+#     # if torch.cuda.is_available():
+#     #     embedding_model = SentenceTransformer(model_name, device='cuda')
+#     # else:
+#     embedding_model = SentenceTransformer(model_name)
+#     print('Load sentence embedding model done')
+#     return embedding_model
+
+
+# def set_index(vector):
+#     if torch.cuda.is_available():
+#         res = faiss.StandardGpuResources()
+#         index = faiss.IndexFlatL2(vector.shape[1])
+#         gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index)
+#         gpu_index_flat.add(vector)
+#         index = gpu_index_flat
+#     else:
+#         index = faiss.IndexFlatL2(vector.shape[1])
+#         index.add(vector)
+#     return index
+
+
+# def get_embeddings(embedding_model, text_list):
+#     return embedding_model.encode(text_list)
+
+
+# def prepare_sentences_vector(encoded_list):
+#     encoded_list = [i.reshape(1, -1) for i in encoded_list]
+#     encoded_list = np.vstack(encoded_list).astype('float32')
+#     encoded_list = normalize(encoded_list)
+#     return encoded_list
+
+
+# def store_embeddings(df, embeddings):
+#     with open('embeddings.pkl', "wb") as fOut:
+#         pickle.dump({'sentences': df['Question'], 'embeddings': embeddings}, fOut, protocol=pickle.HIGHEST_PROTOCOL)
+#     print('Store embeddings done')
+
+
+# def load_embeddings(file_path=EMBEDDINGS_PATH):
+#     with open(file_path, "rb") as fIn:
+#         stored_data = pickle.load(fIn)
+#         stored_sentences = stored_data['sentences']
+#         stored_embeddings = stored_data['embeddings']
+#     print('Load (questions) embeddings done')
+#     return stored_embeddings
+
+
+# def model_pipeline(model, tokenizer, question, similar_context):
+#     inputs = tokenizer(question, similar_context, return_tensors="pt")
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#     answer_start_index = outputs.start_logits.argmax()
+#     answer_end_index = outputs.end_logits.argmax()
+#     predict_answer_tokens = inputs.input_ids[0, answer_start_index: answer_end_index + 1]
+#     Answer = tokenizer.decode(predict_answer_tokens)
+#     return Answer.replace('<unk>','@')
+
+
+# def faiss_search(index, question_vector, k=1):
+#     distances, indices = index.search(question_vector, k)
+#     return distances,indices
+
+# def create_segment_index(vector):
+#     segment_index = faiss.IndexFlatL2(vector.shape[1])
+#     segment_index.add(vector)
+#     return segment_index
+
+
+# def predict_faiss(model, tokenizer, embedding_model, df, question, index):
+#     t = time.time()
+#     question = question.strip()
+#     question_vector = get_embeddings(embedding_model, question)
+#     question_vector = prepare_sentences_vector([question_vector])
+#     distances,indices = faiss_search(index, question_vector)
+#     Answers = [df['Answer'][i] for i in indices[0]]
+#     _time = time.time() - t
+#     output = {
+#         "user_question": question,
+#         "answer": Answers[0],
+#         "totaltime": round(_time, 3),
+#         "score": round(distances[0][0], 4)
+#     }
+#     return output
+
+# def predict(model, tokenizer, embedding_model, df, question, index):
+#     t = time.time()
+#     question = question.strip()
+#     question_vector = get_embeddings(embedding_model, question)
+#     question_vector = prepare_sentences_vector([question_vector])
+#     distances,indices = faiss_search(index, question_vector)
+
+#     # Answer = model_pipeline(model, tokenizer, df['Question'][indices[0][0]], df['Context'][indices[0][0]])
+#     Answer = model_pipeline(model, tokenizer, question, df['Context'][indices[0][0]])
+#     _time = time.time() - t
+#     output = {
+#         "user_question": question,
+#         "answer": Answer,
+#         "totaltime": round(_time, 3),
+#         "distance": round(distances[0][0], 4)
+#     }
+#     return Answer
+
+# def predict_test(model, tokenizer, embedding_model, df, question, index):  # sent_tokenize pythainlp
+#     t = time.time()
+#     question = question.strip()
+#     question_vector = get_embeddings(embedding_model, question)
+#     question_vector = prepare_sentences_vector([question_vector])
+#     distances,indices = faiss_search(index, question_vector)
+
+#     mostSimContext = df['Context'][indices[0][0]]
+#     pattern = r'(?<=\s{10}).*'
+#     matches = re.search(pattern, mostSimContext, flags=re.DOTALL)
+
+#     if matches:
+#         mostSimContext = matches.group(0)
+
+#     mostSimContext = mostSimContext.strip()
+#     mostSimContext = re.sub(r'\s+', ' ', mostSimContext)
+
+#     segments = sent_tokenize(mostSimContext, engine="crfcut")
+
+#     segment_embeddings = get_embeddings(embedding_model, segments)
+#     segment_embeddings = prepare_sentences_vector(segment_embeddings)
+#     segment_index = create_segment_index(segment_embeddings)
+        
+#     _distances,_indices = faiss_search(segment_index, question_vector)
+#     mostSimSegment = segments[_indices[0][0]]
+
+#     Answer = model_pipeline(model, tokenizer,question,mostSimSegment)
+
+#     if len(Answer) <= 2:
+#         Answer = mostSimSegment
+    
+#     # Find the start and end indices of mostSimSegment within mostSimContext
+#     start_index = mostSimContext.find(Answer)
+#     end_index = start_index + len(Answer)
+    
+#     print(f"answer {len(Answer)} => {Answer} || startIndex =>{start_index} || endIndex =>{end_index}")
+#     print(f"mostSimContext{len(mostSimContext)}=>{mostSimContext}\nsegments{len(segments)}=>{segments}\nmostSimSegment{len(mostSimSegment)}=>{mostSimSegment}")
+    
+#     _time = time.time() - t
+#     output = {
+#         "user_question": question,
+#         "answer": df['Answer'][indices[0][0]],
+#         "totaltime": round(_time, 3),
+#         "distance": round(distances[0][0], 4),
+#         "highlight_start": start_index,
+#         "highlight_end": end_index
+#     }
+#     return output
+
+# def highlight_text(text, start_index, end_index):
+#     if start_index < 0:
+#       start_index = 0
+#     if end_index > len(text):
+#       end_index = len(text)
+#     highlighted_text = ""
+#     for i, char in enumerate(text):
+#         if i == start_index:
+#             highlighted_text += "<mark>"
+#         highlighted_text += char
+#         if i == end_index - 1:
+#             highlighted_text += "</mark>"
+#     return highlighted_text
+
+# def chat_interface_before(question, history):
+#     response = predict(model, tokenizer, embedding_model, df, question, index)
+#     return response
+
+# def chat_interface_after(question, history):
+#     response = predict_test(model, tokenizer, embedding_model, df, question, index)
+#     highlighted_answer = highlight_text(response["answer"], response["highlight_start"], response["highlight_end"])
+#     return highlighted_answer
+
+# examples=[
+#             'ขอเลขที่บัญชีของบริษัทหน่อย',
+#             'บริษัทตั้งอยู่ที่ถนนอะไร',
+#             'ขอช่องทางติดตามข่าวสารทาง Line หน่อย',    
+#             'อยากทราบความถี่ในการดึงข้อมูลของ DXT360 ในแต่ละแพลตฟอร์ม',
+#             'อยากทราบความถี่ในการดึงข้อมูลของ DXT360 บน Twitter',
+#             # 'ช่องทางติดตามข่าวสารของเรา',
+#                                 ]
+# demo_before = gr.ChatInterface(fn=chat_interface_before, 
+#                                 examples=examples)
+
+# demo_after = gr.ChatInterface(fn=chat_interface_after, 
+#                               examples=examples)
+
+# interface = gr.TabbedInterface([demo_before, demo_after], ["Before", "After"])
+
+# if __name__ == "__main__":
+#     # Load your model, tokenizer, data, and index here...
+#     df = load_data()
+#     model, tokenizer = load_model('wangchanberta-hyp')
+#     embedding_model = load_embedding_model()
+#     index = set_index(prepare_sentences_vector(load_embeddings(EMBEDDINGS_PATH)))
+#     interface.launch()
+
+
 import random
 import gradio as gr
 import time
@@ -17,224 +258,235 @@ from pythainlp.tokenize import sent_tokenize
 DEFAULT_MODEL = 'wangchanberta-hyp'
 DEFAULT_SENTENCE_EMBEDDING_MODEL = 'intfloat/multilingual-e5-base'
 
-MODEL_DICT = {
-    'wangchanberta': 'Chananchida/wangchanberta-xet_ref-params',
-    'wangchanberta-hyp': 'Chananchida/wangchanberta-xet_hyp-params',
-}
-
-EMBEDDINGS_PATH = 'data/embeddings.pkl'
-DATA_PATH='data/dataset.xlsx'
-
-
-def load_data(path=DATA_PATH):
-    df = pd.read_excel(path, sheet_name='Default')
-    df['Context'] = pd.read_excel(path, sheet_name='mdeberta')['Context']
-    print(len(df))
-    print('Load data done')
-    return df
-
-
-def load_model(model_name=DEFAULT_MODEL):
-    model = AutoModelForQuestionAnswering.from_pretrained(MODEL_DICT[model_name])
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_DICT[model_name])
-    print('Load model done')
-    return model, tokenizer
-
-def load_embedding_model(model_name=DEFAULT_SENTENCE_EMBEDDING_MODEL):
-    # if torch.cuda.is_available():
-    #     embedding_model = SentenceTransformer(model_name, device='cuda')
-    # else:
-    embedding_model = SentenceTransformer(model_name)
-    print('Load sentence embedding model done')
-    return embedding_model
-
-
-def set_index(vector):
-    if torch.cuda.is_available():
-        res = faiss.StandardGpuResources()
-        index = faiss.IndexFlatL2(vector.shape[1])
-        gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index)
-        gpu_index_flat.add(vector)
-        index = gpu_index_flat
-    else:
-        index = faiss.IndexFlatL2(vector.shape[1])
-        index.add(vector)
-    return index
-
-
-def get_embeddings(embedding_model, text_list):
-    return embedding_model.encode(text_list)
-
-
-def prepare_sentences_vector(encoded_list):
-    encoded_list = [i.reshape(1, -1) for i in encoded_list]
-    encoded_list = np.vstack(encoded_list).astype('float32')
-    encoded_list = normalize(encoded_list)
-    return encoded_list
-
-
-def store_embeddings(df, embeddings):
-    with open('embeddings.pkl', "wb") as fOut:
-        pickle.dump({'sentences': df['Question'], 'embeddings': embeddings}, fOut, protocol=pickle.HIGHEST_PROTOCOL)
-    print('Store embeddings done')
-
-
-def load_embeddings(file_path=EMBEDDINGS_PATH):
-    with open(file_path, "rb") as fIn:
-        stored_data = pickle.load(fIn)
-        stored_sentences = stored_data['sentences']
-        stored_embeddings = stored_data['embeddings']
-    print('Load (questions) embeddings done')
-    return stored_embeddings
-
-
-def model_pipeline(model, tokenizer, question, similar_context):
-    inputs = tokenizer(question, similar_context, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    answer_start_index = outputs.start_logits.argmax()
-    answer_end_index = outputs.end_logits.argmax()
-    predict_answer_tokens = inputs.input_ids[0, answer_start_index: answer_end_index + 1]
-    Answer = tokenizer.decode(predict_answer_tokens)
-    return Answer.replace('<unk>','@')
-
-
-def faiss_search(index, question_vector, k=1):
-    distances, indices = index.search(question_vector, k)
-    return distances,indices
-
-def create_segment_index(vector):
-    segment_index = faiss.IndexFlatL2(vector.shape[1])
-    segment_index.add(vector)
-    return segment_index
-
-
-def predict_faiss(model, tokenizer, embedding_model, df, question, index):
-    t = time.time()
-    question = question.strip()
-    question_vector = get_embeddings(embedding_model, question)
-    question_vector = prepare_sentences_vector([question_vector])
-    distances,indices = faiss_search(index, question_vector)
-    Answers = [df['Answer'][i] for i in indices[0]]
-    _time = time.time() - t
-    output = {
-        "user_question": question,
-        "answer": Answers[0],
-        "totaltime": round(_time, 3),
-        "score": round(distances[0][0], 4)
-    }
-    return output
-
-def predict(model, tokenizer, embedding_model, df, question, index):
-    t = time.time()
-    question = question.strip()
-    question_vector = get_embeddings(embedding_model, question)
-    question_vector = prepare_sentences_vector([question_vector])
-    distances,indices = faiss_search(index, question_vector)
-
-    # Answer = model_pipeline(model, tokenizer, df['Question'][indices[0][0]], df['Context'][indices[0][0]])
-    Answer = model_pipeline(model, tokenizer, question, df['Context'][indices[0][0]])
-    _time = time.time() - t
-    output = {
-        "user_question": question,
-        "answer": Answer,
-        "totaltime": round(_time, 3),
-        "distance": round(distances[0][0], 4)
-    }
-    return Answer
-
-def predict_test(model, tokenizer, embedding_model, df, question, index):  # sent_tokenize pythainlp
-    t = time.time()
-    question = question.strip()
-    question_vector = get_embeddings(embedding_model, question)
-    question_vector = prepare_sentences_vector([question_vector])
-    distances,indices = faiss_search(index, question_vector)
-
-    mostSimContext = df['Context'][indices[0][0]]
-    pattern = r'(?<=\s{10}).*'
-    matches = re.search(pattern, mostSimContext, flags=re.DOTALL)
-
-    if matches:
-        mostSimContext = matches.group(0)
-
-    mostSimContext = mostSimContext.strip()
-    mostSimContext = re.sub(r'\s+', ' ', mostSimContext)
-
-    segments = sent_tokenize(mostSimContext, engine="crfcut")
-
-    segment_embeddings = get_embeddings(embedding_model, segments)
-    segment_embeddings = prepare_sentences_vector(segment_embeddings)
-    segment_index = create_segment_index(segment_embeddings)
+MODEL_DICT = 'Chananchida/wangchanberta-xet_hyp-params'
+
+
+EMBEDDINGS_PATH = '/content/dataset.xlsx'
+DATA_PATH='/content/embeddings.pkl'
+
+class ChatBot:   
+    SHEET_NAME_MDEBERTA = 'mdeberta'
+    SHEET_NAME_DEFAULT = 'Default'
+    UNKNOWN_ANSWERS = ["กรุณาลงรายระเอียดมากกว่านี้ได้มั้ยคะ", "ขอโทษค่ะลูกค้า ดิฉันไม่ทราบจริง ๆ"]
+
+    def __init__(self, df_path=None, model_path=None, tokenizer_path=None, embedding_model_name=None, embeddingsPath=None):
+        self.df = None
+        self.model = None
+        self.tokenizer = None
+        self.embedding_model = None
+        self.index = None
+        self.k = 5
+        if all(arg is not None for arg in (df_path, model_path, tokenizer_path, embedding_model_name, embeddingsPath)):
+            self.set_df(df_path)
+            self.set_model(model_path)
+            self.set_tokenizer(tokenizer_path)
+            self.set_embedding_model(embedding_model_name)
+            sentences_vector = self.load_embeddings(embeddingsPath)
+            repared_vector = self.prepare_sentences_vector(sentences_vector)
+            self.set_index(repared_vector)
+            
+    def set_index(self, vector):
+        if torch.cuda.is_available():  # Check if GPU is available
+            res = faiss.StandardGpuResources()
+            index = faiss.IndexFlatL2(vector.shape[1])
+            gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index)
+            gpu_index_flat.add(vector)
+            self.index = gpu_index_flat
+        else:  # If GPU is not available, use CPU-based Faiss index
+            self.index = faiss.IndexFlatL2(vector.shape[1])
+            self.index.add(vector)
+        return self.index 
+
+    def set_df(self, path):
+        self.df = pd.read_excel(path, sheet_name=self.SHEET_NAME_DEFAULT)
+        self.df.rename(columns={'Response': 'Answer'}, inplace=True)
+        self.df['Context'] = pd.read_excel(path, self.SHEET_NAME_MDEBERTA)['Context']
+
+    def set_model(self, model):
+        self.model = AutoModelForQuestionAnswering.from_pretrained(model)
+
+    def set_tokenizer(self, tokenizer):
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+
+    def set_embedding_model(self, model):
+        self.embedding_model = SentenceTransformer(model)
+
+    def set_k(self, k_value):
+        self.k = k_value
+
+    def get_df(self):
+        return self.df
+
+    def get_model(self):
+        return self.model
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def get_embedding_model(self):
+        return self.embedding_model
+
+    def get_index(self):
+        return self.index
+
+    def get_k(self):
+        return self.k
+
+    def get_embeddings(self, text_list):
+        return self.embedding_model.encode(text_list)
+
+    def prepare_sentences_vector(self, encoded_list):
+        encoded_list = [i.reshape(1, -1) for i in encoded_list]
+        encoded_list = np.vstack(encoded_list).astype('float32')
+        encoded_list = normalize(encoded_list)
+        return encoded_list
+
+    def load_embeddings(self, file_path):
+        with open(file_path, "rb") as fIn:
+            stored_data = pickle.load(fIn)
+            stored_sentences = stored_data['sentences']
+            stored_embeddings = stored_data['embeddings']
+        return stored_embeddings
+
+    def model_pipeline(self, question, similar_context):
+        inputs = self.tokenizer(question, similar_context, return_tensors="pt")
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        answer_start_index = outputs.start_logits.argmax()
+        answer_end_index = outputs.end_logits.argmax()
+        predict_answer_tokens = inputs.input_ids[0, answer_start_index: answer_end_index + 1]
+        Answer = self.tokenizer.decode(predict_answer_tokens)
+        return Answer.replace('<unk>','@')
+
+    def faiss_search(self, index, question_vector):
+        if index is None:
+            raise ValueError("Index has not been initialized.")
+        distances, indices = index.search(question_vector, self.k)
+        similar_questions = [self.df['Question'][indices[0][i]] for i in range(self.k)]
+        similar_contexts = [self.df['Context'][indices[0][i]] for i in range(self.k)]
+        return similar_questions, similar_contexts, distances, indices
+    
+    def faiss_segment_search(self, index, question_vector, x=1):
+        if index is None:
+            raise ValueError("Index has not been initialized.")
+        distances, indices = index.search(question_vector, x)
+        return distances, indices
+    
+    def create_segment_index(self, vector):
+        segment_index = faiss.IndexFlatL2(vector.shape[1])
+        segment_index.add(vector)
+        return segment_index
+
+    def predict_test(self, question):
+        list_context_for_show = []
+        list_distance_for_show = []
+        list_similar_question = []
+
+        question = question.strip()
+        question_vector = self.get_embeddings([question])
+        question_vector = self.prepare_sentences_vector([question_vector])
+        similar_questions, similar_contexts, distances, indices = self.faiss_search(self.index, question_vector)
+
+        
+        mostSimContext = similar_contexts[0]
+        pattern = r'(?<=\s{10}).*'
+        matches = re.search(pattern, mostSimContext, flags=re.DOTALL)
+        
+        if matches:
+            mostSimContext = matches.group(0)
+            
+        mostSimContext = mostSimContext.strip()
+        mostSimContext = re.sub(r'\s+', ' ', mostSimContext)
         
-    _distances,_indices = faiss_search(segment_index, question_vector)
-    mostSimSegment = segments[_indices[0][0]]
+        segments = sent_tokenize(mostSimContext, engine="crfcut")
 
-    Answer = model_pipeline(model, tokenizer,question,mostSimSegment)
+        segment_embeddings = self.get_embeddings(segments)
+        segment_embeddings = self.prepare_sentences_vector(segment_embeddings)
+        segment_index = self.create_segment_index(segment_embeddings)
 
-    if len(Answer) <= 2:
-        Answer = mostSimSegment
+        _distances, _indices = self.faiss_segment_search(segment_index, question_vector)
     
-    # Find the start and end indices of mostSimSegment within mostSimContext
-    start_index = mostSimContext.find(Answer)
-    end_index = start_index + len(Answer)
-    
-    print(f"answer {len(Answer)} => {Answer} || startIndex =>{start_index} || endIndex =>{end_index}")
-    print(f"mostSimContext{len(mostSimContext)}=>{mostSimContext}\nsegments{len(segments)}=>{segments}\nmostSimSegment{len(mostSimSegment)}=>{mostSimSegment}")
-    
-    _time = time.time() - t
-    output = {
-        "user_question": question,
-        "answer": df['Answer'][indices[0][0]],
-        "totaltime": round(_time, 3),
-        "distance": round(distances[0][0], 4),
-        "highlight_start": start_index,
-        "highlight_end": end_index
-    }
-    return output
-
-def highlight_text(text, start_index, end_index):
-    if start_index < 0:
-      start_index = 0
-    if end_index > len(text):
-      end_index = len(text)
-    highlighted_text = ""
-    for i, char in enumerate(text):
-        if i == start_index:
-            highlighted_text += "<mark>"
-        highlighted_text += char
-        if i == end_index - 1:
-            highlighted_text += "</mark>"
-    return highlighted_text
-
-def chat_interface_before(question, history):
-    response = predict(model, tokenizer, embedding_model, df, question, index)
-    return response
-
-def chat_interface_after(question, history):
-    response = predict_test(model, tokenizer, embedding_model, df, question, index)
-    highlighted_answer = highlight_text(response["answer"], response["highlight_start"], response["highlight_end"])
-    return highlighted_answer
-
-examples=[
-            'ขอเลขที่บัญชีของบริษัทหน่อย',
-            'บริษัทตั้งอยู่ที่ถนนอะไร',
-            'ขอช่องทางติดตามข่าวสารทาง Line หน่อย',    
-            'อยากทราบความถี่ในการดึงข้อมูลของ DXT360 ในแต่ละแพลตฟอร์ม',
-            'อยากทราบความถี่ในการดึงข้อมูลของ DXT360 บน Twitter',
-            # 'ช่องทางติดตามข่าวสารของเรา',
-                                ]
-demo_before = gr.ChatInterface(fn=chat_interface_before, 
-                                examples=examples)
-
-demo_after = gr.ChatInterface(fn=chat_interface_after, 
-                              examples=examples)
-
-interface = gr.TabbedInterface([demo_before, demo_after], ["Before", "After"])
+        mostSimSegment = segments[_indices[0][0]]
 
+        print(f"_indices => {_indices[0][0]}")
+        answer = self.model_pipeline(question, mostSimSegment)
+        
+        if len(answer) <= 2:
+            answer = mostSimSegment
+        
+        start_index = mostSimContext.find(answer)
+        end_index = start_index + len(answer)
+
+        print(f"mostSimContext {len(mostSimContext)} =>{mostSimContext}\nsegments {len(segments)} =>{segments}\nmostSimSegment {len(mostSimSegment)} =>{mostSimSegment}")
+        print(f"answer {len(answer)} => {answer} || startIndex =>{start_index} || endIndex =>{end_index}")
+
+        for i in range(min(5, self.k)):
+            index = indices[0][i]
+            similar_question = similar_questions[i]
+            similar_context = similar_contexts[i]
+
+            list_similar_question.append(similar_question)
+            list_context_for_show.append(similar_context)
+            list_distance_for_show.append(str(1 - distances[0][i]))
+
+        distance = list_distance_for_show[0]
+
+        if float(distance) < 0.5:
+            answer = random.choice(self.UNKNOWN_ANSWERS)
+
+        output = {
+            "user_question": question,
+            "answer": self.df['Answer'][indices[0][0]],
+            "distance": distance,
+            "highlight_start": start_index,
+            "highlight_end": end_index,
+            "list_context": list_context_for_show,
+            "list_distance": list_distance_for_show
+        }
+        return output
+
+    def highlight_text(self, text, start_index, end_index):
+        if start_index < 0:
+            start_index = 0
+        if end_index > len(text):
+            end_index = len(text)
+        highlighted_text = ""
+        for i, char in enumerate(text):
+            if i == start_index:
+                highlighted_text += "<mark>"
+            highlighted_text += char
+            if i == end_index - 1:
+                highlighted_text += "</mark>"
+        return highlighted_text
+
+    def chat_interface_before(self, question, history):
+        response = self.predict(question)
+        return response
+
+    def chat_interface_after(self, question, history):
+        response = self.predict_test(question)
+        highlighted_answer = self.highlight_text(response["answer"], response["highlight_start"], response["highlight_end"])
+        return highlighted_answer
+  
 if __name__ == "__main__":
-    # Load your model, tokenizer, data, and index here...
-    df = load_data()
-    model, tokenizer = load_model('wangchanberta-hyp')
-    embedding_model = load_embedding_model()
-    index = set_index(prepare_sentences_vector(load_embeddings(EMBEDDINGS_PATH)))
-    interface.launch()
+  bot = ChatBot(df_path=DATA_PATH, model_path=MODEL_DICT, tokenizer_path=MODEL_DICT, embedding_model_name=DEFAULT_SENTENCE_EMBEDDING_MODEL, embeddingsPath=EMBEDDINGS_PATH)
+  # bot.load_data()
+  # bot.load_model()
+  # bot.load_embedding_model()
+  # embeddings = bot.load_embeddings(EMBEDDINGS_PATH)
+  # bot.set_index(bot.prepare_sentences_vector(embeddings))
+
+  examples = [
+      'ขอเลขที่บัญชีของบริษัทหน่อย',
+      'บริษัทตั้งอยู่ที่ถนนอะไร',
+      'ขอช่องทางติดตามข่าวสารทาง Line หน่อย',
+      'อยากทราบความถี่ในการดึงข้อมูลของ DXT360 ในแต่ละแพลตฟอร์ม',
+      'อยากทราบความถี่ในการดึงข้อมูลของ DXT360 บน Twitter',
+      # 'ช่องทางติดตามข่าวสารของเรา',
+  ]
+
+  demo_before = gr.ChatInterface(fn=bot.chat_interface_before, examples=examples)
+  demo_after = gr.ChatInterface(fn=bot.chat_interface_after, examples=examples)
+
+  interface = gr.TabbedInterface([demo_before, demo_after], ["Before", "After"])
+  interface.launch()