add retrieval for decoded and latent

Author: haok1402

script/acquire.sh

@@ -1,9 +1,17 @@
 #!/usr/bin/bash
 
-# Acquire a compute node with 32 CPUs, 96GB RAM, and 1 A6000 GPU.
+# Validate the arguments
+if [ -z "$1" ]; then
+    echo "Usage: $0 <num_gpu>"
+    exit 1
+fi
+
+NUM_GPU=$1
+
+# Acquire a compute node with 32 CPUs, 96GB RAM, and specified number of GPUs.
 # The node will be acquired for 7 days, and the session will be interactive.
 # Please run this script with tmux to avoid losing the session.
 srun \
-    --partition=long --time=07-00:00:00 \
-    --cpus-per-task=32 --mem=96GB --gres=gpu:A6000:1 \
-    --pty bash
+        --partition=long --time=07-00:00:00 \
+        --cpus-per-task=32 --mem=96GB --gres=gpu:A6000:$NUM_GPU \
+        --pty bash

source/interpret/retrieval/reconstruct.py source/interpret/retrieval/decoded.py

@@ -33,7 +33,7 @@ def main(dataset: Dataset, embedding: Type[Embedding], version: str):
     )
 
     # define where to save results
-    saveBase = Path(workspace, "reconstruct")
+    saveBase = Path(workspace, "decoded")
     saveBase.mkdir(mode=0o770, parents=True, exist_ok=True)
     qresFile = Path(saveBase, f"{version}.qres")
     evalFile = Path(saveBase, f"{version}.eval")
@@ -54,7 +54,7 @@ def main(dataset: Dataset, embedding: Type[Embedding], version: str):
                 D, I = gpuIndex.search(bQrys, 100)
                 for qid, sims, dnos in zip(bQids, D, I):
                     for s, d in zip(sims, dnos):
-                        f.write(f"{qid}\tQ0\t{dids[d]}\t0\t{s}\tReconstruct\n")
+                        f.write(f"{qid}\tQ0\t{dids[d]}\t0\t{s}\tDecoded\n")
                 f.flush()
                 p.update(t, advance=bQrys.shape[0])
         p.remove_task(t)

source/interpret/retrieval/latent.py

@@ -14,7 +14,6 @@
 from source.dataset import MsMarcoDataset
 from source.embedding import BgeBaseEmbedding
 from source.interface import Dataset, Embedding
-from source.interpret import esUser, esAuth, esCert
 
 
 def main(
@@ -60,10 +59,7 @@ def main(
 
     # create connection to elastic search
     es = Elasticsearch(
-        hosts=[{"host": esHost, "port": esPort}],
-        http_auth=(esUser, esAuth),
-        ca_certs=esCert,
-        scheme="https",
+        hosts=[{"host": esHost, "port": esPort, "scheme": "http"}],
     )
     es.indices.create(
         index=f"{version}.latent".lower(),
@@ -163,7 +159,7 @@ def main(
     parser.add_argument("version", type=str)
     parser.add_argument("--dataset", type=str, default="MsMarco", choices=["MsMarco"])
     parser.add_argument("--embedding", type=str, default="BgeBase", choices=["BgeBase"])
-    parser.add_argument("--esHost", type=str, default="172.16.1.166")
+    parser.add_argument("--esHost", type=str, default="localhost")
     parser.add_argument("--esPort", type=int, default=9200)
     args = parser.parse_args()