Update app.py

app.py

@@ -9,7 +9,7 @@ from torch import Tensor
 from transformers import AutoTokenizer, AutoModel
 from sklearn.metrics.pairwise import cosine_similarity
 
-import traceback
+import re
 
 
 def average_pool(last_hidden_states: Tensor,
@@ -65,7 +65,8 @@ with gr.Blocks() as demo:
                     if str(idx) == n.split(".")[0]:
                         df_idx = int(n.split(" // ")[-1])
                         print(df_idx)
-                        article = df.iloc[df_idx]['text'].replace("==", "\n\n==")
+                        article = df.iloc[df_idx]['text']
+                        article = re.sub(r'(===?=?[A-Z ].+?===?=?)', r'\n\n\1\n', article)
                         chat_history.append((message, f"contents of {n}:\n{article}"))
                         return "", chat_history
         print("nothing found")