Update app.py

app.py

@@ -1,6 +1,3 @@
-# -*- coding: utf-8 -*-
-
-#@title scirpts
 import time
 import numpy as np
 import pandas as pd
@@ -8,8 +5,7 @@ import torch
 import faiss
 from sklearn.preprocessing import normalize
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering
-from sentence_transformers import SentenceTransformer,util
-from pythainlp import Tokenizer
+from sentence_transformers import SentenceTransformer
 import pickle
 import gradio as gr
 
@@ -17,7 +13,7 @@ print(torch.cuda.is_available())
 
 __all__ = [
     "mdeberta",
-    "wangchanberta-hyp", # Best model
+    "wangchanberta-hyp",  # Best model
 ]
 
 predict_method = [
@@ -27,8 +23,8 @@ predict_method = [
     "semanticSearchWithModel",
 ]
 
-DEFAULT_MODEL='wangchanberta-hyp'
-DEFAULT_SENTENCE_EMBEDDING_MODEL='intfloat/multilingual-e5-base'
+DEFAULT_MODEL = 'wangchanberta-hyp'
+DEFAULT_SENTENCE_EMBEDDING_MODEL = 'intfloat/multilingual-e5-base'
 
 MODEL_DICT = {
     'wangchanberta': 'Chananchida/wangchanberta-th-wiki-qa_ref-params',
@@ -37,8 +33,8 @@ MODEL_DICT = {
     'mdeberta-hyp': 'Chananchida/mdeberta-v3-th-wiki-qa_hyp-params',
 }
 
-DATA_PATH='models/dataset.xlsx'
-EMBEDDINGS_PATH='models/embeddings.pkl'
+DATA_PATH = 'models/dataset.xlsx'
+EMBEDDINGS_PATH = 'models/embeddings.pkl'
 
 
 class ChatbotModel:
@@ -50,12 +46,12 @@ class ChatbotModel:
         self._chatbot.set_vectors()
         self._chatbot.set_index()
 
-
     def chat(self, question):
         return self._chatbot.answer_question(question)
 
-    def eval(self,model,predict_method):
-        return self._chatbot.eval(model_name=model,predict_method=predict_method)
+    def eval(self, model, predict_method):
+        return self._chatbot.eval(model_name=model, predict_method=predict_method)
+
 
 class Chatbot:
     def __init__(self):
@@ -73,31 +69,29 @@ class Chatbot:
     def load_data(self, path: str = DATA_PATH):
         self.df = pd.read_excel(path, sheet_name='Default')
         self.df['Context'] = pd.read_excel(path, sheet_name='mdeberta')['Context']
-        # print('Load data done')
 
     def load_model(self, model_name: str = DEFAULT_MODEL):
         self.model = AutoModelForQuestionAnswering.from_pretrained(MODEL_DICT[model_name])
         self.tokenizer = AutoTokenizer.from_pretrained(MODEL_DICT[model_name])
         self.model_name = model_name
-        # print('Load model done')
 
     def load_embedding_model(self, model_name: str = DEFAULT_SENTENCE_EMBEDDING_MODEL):
-        if torch.cuda.is_available():  # Check if GPU is available
-            self.embedding_model = SentenceTransformer(model_name, device='cpu')
-        else: self.embedding_model = SentenceTransformer(model_name)
-        # print('Load sentence embedding model done')
+        if torch.cuda.is_available():
+            self.embedding_model = SentenceTransformer(model_name, device='cuda')
+        else:
+            self.embedding_model = SentenceTransformer(model_name)
 
     def set_vectors(self):
         self.vectors = self.prepare_sentences_vector(self.load_embeddings(EMBEDDINGS_PATH))
 
     def set_index(self):
-        if torch.cuda.is_available():  # Check if GPU is available
+        if torch.cuda.is_available():
             res = faiss.StandardGpuResources()
             self.index = faiss.IndexFlatL2(self.vectors.shape[1])
             gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, self.index)
             gpu_index_flat.add(self.vectors)
             self.index = gpu_index_flat
-        else:  # If GPU is not available, use CPU-based Faiss index
+        else:
             self.index = faiss.IndexFlatL2(self.vectors.shape[1])
             self.index.add(self.vectors)
 
@@ -110,18 +104,15 @@ class Chatbot:
         encoded_list = normalize(encoded_list)
         return encoded_list
 
-
     def store_embeddings(self, embeddings):
         with open('models/embeddings.pkl', "wb") as fOut:
             pickle.dump({'sentences': self.df['Question'], 'embeddings': embeddings}, fOut, protocol=pickle.HIGHEST_PROTOCOL)
-        print('Store embeddings done')
 
     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']
-        print('Load (questions) embeddings done')
         return stored_embeddings
 
     def model_pipeline(self, question, similar_context):
@@ -140,25 +131,24 @@ class Chatbot:
         similar_contexts = [self.df['Context'][indices[0][i]] for i in range(self.k)]
         return similar_questions, similar_contexts, distances, indices
 
-
-    def predict(self,message):
+    def predict(self, message):
         message = message.strip()
         question_vector = self.get_embeddings(message)
-        question_vector=self.prepare_sentences_vector([question_vector])
-        similar_questions, similar_contexts, distances,indices = self.faiss_search(question_vector)
+        question_vector = self.prepare_sentences_vector([question_vector])
+        similar_questions, similar_contexts, distances, indices = self.faiss_search(question_vector)
         Answer = self.model_pipeline(message, similar_contexts)
         start_index = similar_contexts.find(Answer)
         end_index = start_index + len(Answer)
-        _time = time.time() - t
         output = {
             "user_question": message,
-            "answer": df['Answer'][indices[0][0]],
-            "totaltime": round(_time, 3),
+            "answer": self.df['Answer'][indices[0][0]],
             "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
@@ -166,21 +156,21 @@ def highlight_text(text, start_index, end_index):
         end_index = len(text)
     highlighted_text = ""
     for i, char in enumerate(text):
-         if i == start_index:
+        if i == start_index:
             highlighted_text += "<mark>"
         highlighted_text += char
         if i == end_index - 1:
             highlighted_text += "</mark>"
     return highlighted_text
 
-    
-"""#Gradio"""
-if __name__ == "__main__":        
+
+if __name__ == "__main__":
     bot = ChatbotModel()
+    
     def chat_interface(question, history):
-        response = bot._chatbot.predict(model, tokenizer, embedding_model, df, question, index)
+        response = bot._chatbot.predict(question)
         highlighted_answer = highlight_text(response["answer"], response["highlight_start"], response["highlight_end"])
         return highlighted_answer
-    # EXAMPLE = ["หลิน ไห่เฟิง มีชื่อเรียกอีกชื่อว่าอะไร" , "ใครเป็นผู้ตั้งสภาเศรษฐกิจโลกขึ้นในปี พ.ศ. 2514 โดยทุกปีจะมีการประชุมที่ประเทศสวิตเซอร์แลนด์", "โปรดิวเซอร์ของอัลบั้มตลอดกาล ของวงคีรีบูนคือใคร", "สกุลเดิมของหม่อมครูนุ่ม นวรัตน ณ อยุธยา คืออะไร"]
-    demo = gr.ChatInterface(fn=chat_interface, title="CE66-04_Thai Question Answering System by using Deep Learning")
+    
+    demo = gr.Interface(fn=chat_interface, title="Thai Question Answering System", inputs="text", outputs="html")
     demo.launch()