Update app.py

app.py

@@ -27,10 +27,10 @@ DEFAULT_MODEL = 'wangchanberta-hyp'
 DEFAULT_SENTENCE_EMBEDDING_MODEL = 'intfloat/multilingual-e5-base'
 
 MODEL_DICT = {
-    'wangchanberta': 'Chananchida/wangchanberta-th-wiki-qa_ref-params',
-    'wangchanberta-hyp': 'Chananchida/wangchanberta-th-wiki-qa_hyp-params',
-    'mdeberta': 'Chananchida/mdeberta-v3-th-wiki-qa_ref-params',
-    'mdeberta-hyp': 'Chananchida/mdeberta-v3-th-wiki-qa_hyp-params',
+    'wangchanberta': 'powerpuf-bot/wangchanberta-th-wiki-qa_ref-params',
+    'wangchanberta-hyp': 'powerpuf-bot/wangchanberta-th-wiki-qa_hyp-params',
+    'mdeberta': 'powerpuf-bot/mdeberta-v3-th-wiki-qa_ref-params',
+    'mdeberta-hyp': 'powerpuf-bot/mdeberta-v3-th-wiki-qa_hyp-params',
 }
 
 DATA_PATH = 'models/dataset.xlsx'
@@ -138,40 +138,26 @@ class Chatbot:
         similar_questions, similar_contexts, distances, indices = self.faiss_search(question_vector)
         context = similar_contexts[0]
         Answer = self.model_pipeline(str(message), context)
-        start_index = context.find(Answer)
-        end_index = start_index + len(Answer)
-        output = {
-            "user_question": message,
-            "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
-    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
-
+        # start_index = context.find(Answer)
+        # end_index = start_index + len(Answer)
+        # output = {
+        #     "user_question": message,
+        #     "answer": self.df['Answer'][indices[0][0]],
+        #     "distance": round(distances[0][0], 4),
+        #     "highlight_start": start_index,
+        #     "highlight_end": end_index
+        # }
+        # return output
+        return Answer
 
 if __name__ == "__main__":
     bot = ChatbotModel()
     
     def chat_interface(question, history):
         response = bot._chatbot.predict(question)
-        highlighted_answer = highlight_text(response["answer"], response["highlight_start"], response["highlight_end"])
-        return highlighted_answer
+        # highlighted_answer = highlight_text(response["answer"], response["highlight_start"], response["highlight_end"])
+        # return highlighted_answer
+        return response
 
     EXAMPLE = ["หลิน ไห่เฟิง มีชื่อเรียกอีกชื่อว่าอะไร" , "ใครเป็นผู้ตั้งสภาเศรษฐกิจโลกขึ้นในปี พ.ศ. 2514 โดยทุกปีจะมีการประชุมที่ประเทศสวิตเซอร์แลนด์", "โปรดิวเซอร์ของอัลบั้มตลอดกาล ของวงคีรีบูนคือใคร", "สกุลเดิมของหม่อมครูนุ่ม นวรัตน ณ อยุธยา คืออะไร"]