Update app.py

app.py

@@ -28,7 +28,7 @@ with gr.Blocks() as demo:
     msg = gr.Textbox(label="simple wikipedia semantic search query", placeholder="for example, \"medieval battles\"")
     clear = gr.ClearButton([msg, chatbot])
 
-    def respond(message, chat_history):
+    def _search(message, chat_history):
         batch_dict = tokenizer(["query: " + message], max_length=512, padding=True, truncation=True, return_tensors='pt')
     
         outputs = model(**batch_dict)
@@ -49,11 +49,28 @@ with gr.Blocks() as demo:
         # Get corresponding 'text' for top k similar points
         top_k_text = df['title'].iloc[top_k_idx].tolist()
         
-        bot_message = "\n".join(f"{i+1}. {top_k_text[i]}" for i in range(len(top_k_text)))
+        bot_message = "\n".join(f"{i+1}. {top_k_text[i]} // {top_k_idx[i]}" for i in range(len(top_k_text)))
     
-        chat_history.append((message, bot_message))
+        chat_history.append((message, "results:\n" + bot_message))
         return "", chat_history
 
+    def _retrieve(message, chat_history):
+        idx = int(message)
+        for _, m in chat_history[::-1]:
+            if m.startswith("results:\n"):
+                for n in m.split("\n")[1:]:
+                    if str(idx) == n.split(".")[0]:
+                        df_idx = int(n.split(" // ")[-1])
+                        chat_history.append((message, f"contents of {n}:\n{df[df_idx]}"))
+                        return "", chat_history
+    
+    def respond(message, chat_history):
+        try:
+            int(message)
+            return _retrieve(message, chat_history)
+        except:
+            return _search(message, chat_history)
+
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
 
 demo.launch()