add vote

app.py

@@ -40,7 +40,7 @@ def calculate_md5_from_binary(binary_data):
     hash_md5.update(binary_data)
     return hash_md5.hexdigest()
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=100)
 def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
     global model, tokenizer
 
@@ -84,6 +84,7 @@ def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
     
     return knowledge_base_name
 
+# @spaces.GPU
 def retrieve_gradio(knowledge_base: str, query: str, topk: int):
     global model, tokenizer
 
@@ -104,12 +105,6 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
         query_rep = model(text=[query_with_instruction], image=[None], tokenizer=tokenizer).reps.squeeze(0).cpu()
 
     query_md5 = hashlib.md5(query.encode()).hexdigest()
-    with open(os.path.join(target_cache_dir, f"q-{query_md5}.json"), 'w') as f:
-        f.write(json.dumps(
-            {
-                "query": query
-            }, indent=4, ensure_ascii=False
-        ))
 
     doc_reps_cat = torch.stack([torch.Tensor(i) for i in doc_reps], dim=0)
 
@@ -125,9 +120,57 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
 
     images_topk = [Image.open(os.path.join(target_cache_dir, f"{md5s[idx]}.png")) for idx in topk_doc_ids_np]
 
+    with open(os.path.join(target_cache_dir, f"q-{query_md5}.json"), 'w') as f:
+        f.write(json.dumps(
+            {
+                "knowledge_base": knowledge_base,
+                "query": query,
+                "retrived_docs": [os.path.join(target_cache_dir, f"{md5s[idx]}.png") for idx in topk_doc_ids_np]
+            }, indent=4, ensure_ascii=False
+        ))
+
     return images_topk
 
 
+def upvote(knowledge_base, query):
+    global model, tokenizer
+
+    target_cache_dir = os.path.join(cache_dir, knowledge_base)
+
+    query_md5 = hashlib.md5(query.encode()).hexdigest()
+
+    with open(os.path.join(target_cache_dir, f"q-{query_md5}.json"), 'r') as f:
+        data = json.loads(f.read())
+
+    data["user_preference"] = "upvote"
+
+    with open(os.path.join(target_cache_dir, f"q-{query_md5}-withpref.json"), 'w') as f:
+        f.write(json.dumps(data, indent=4, ensure_ascii=False))
+
+    print("up", os.path.join(target_cache_dir, f"q-{query_md5}-withpref.json"))
+
+    return
+
+
+def downvote(knowledge_base, query):
+    global model, tokenizer
+
+    target_cache_dir = os.path.join(cache_dir, knowledge_base)
+
+    query_md5 = hashlib.md5(query.encode()).hexdigest()
+
+    with open(os.path.join(target_cache_dir, f"q-{query_md5}.json"), 'r') as f:
+        data = json.loads(f.read())
+
+    data["user_preference"] = "downvote"
+
+    with open(os.path.join(target_cache_dir, f"q-{query_md5}-withpref.json"), 'w') as f:
+        f.write(json.dumps(data, indent=4, ensure_ascii=False))
+
+    print("down", os.path.join(target_cache_dir, f"q-{query_md5}-withpref.json"))
+    return
+
+
 device = 'cuda' 
 model_path = 'RhapsodyAI/minicpm-visual-embedding-v0' # replace with your local model path
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
@@ -136,9 +179,10 @@ model.to(device)
 
 
 with gr.Blocks() as app:
-    gr.Markdown("# Memex: OCR-free Visual Document Retrieval @RhapsodyAI [RhapsodyAI/minicpm-visual-embedding-v0](https://huggingface.co/RhapsodyAI/minicpm-visual-embedding-v0)")
+    gr.Markdown("# Memex: OCR-free Visual Document Retrieval @RhapsodyAI")
 
-    gr.Markdown("Only PDF file is available now. You are suggested to uplaod <= 50 pages per PDF, otherwise you will reach GPU time limit.")
+    gr.Markdown("- We open-sourced our model at [RhapsodyAI/minicpm-visual-embedding-v0](https://huggingface.co/RhapsodyAI/minicpm-visual-embedding-v0)")
+    gr.Markdown("- Currently we support PDF document with less than 50 pages, PDF over 50 pages will reach GPU time limit.")
     
     with gr.Row():
         file_input = gr.File(type="binary", label="Upload PDF")
@@ -148,16 +192,23 @@ with gr.Blocks() as app:
     process_button.click(add_pdf_gradio, inputs=[file_input], outputs=file_result)
 
     with gr.Row():
-        kb_id_input = gr.Text(label="Your Knowledge Base ID")
+        kb_id_input = gr.Text(label="Your Knowledge Base ID (paste your Knowledge Base ID here:)")
         query_input = gr.Text(label="Your Queston")
         topk_input = inputs=gr.Number(value=1, minimum=1, maximum=5, step=1, label="Top K")
         retrieve_button = gr.Button("Retrieve")
     
+    with gr.Row():
+        downvote_button = gr.Button("🤣Downvote")
+        upvote_button = gr.Button("🤗Upvote")
+    
     with gr.Row():
         images_output = gr.Gallery(label="Retrieved Pages")
     
     retrieve_button.click(retrieve_gradio, inputs=[kb_id_input, query_input, topk_input], outputs=images_output)
 
+    upvote_button.click(upvote, inputs=[kb_id_input, query_input], outputs=None)
+    downvote_button.click(downvote, inputs=[kb_id_input, query_input], outputs=None)
+
     gr.Markdown("By using this demo, you agree to share your use data with us for research purpose, to help improve user experience.")
 
 app.launch()