update: retrieved chunks data structure

medrag_multi_modal/retrieval/bm25s_retrieval.py

@@ -175,7 +175,7 @@ class BM25sRetriever(weave.Model):
             results.documents.flatten().tolist(),
             results.scores.flatten().tolist(),
         ):
-            retrieved_chunks.append({"chunk": chunk, "score": score})
+            retrieved_chunks.append({**chunk, **{"score": score}})
         return retrieved_chunks
 
     @weave.op()

medrag_multi_modal/retrieval/contriever_retrieval.py

@@ -192,8 +192,8 @@ class ContrieverRetriever(weave.Model):
         for score in scores:
             retrieved_chunks.append(
                 {
-                    "chunk": self._chunk_dataset[score["original_index"]],
-                    "score": score["item"],
+                    **self._chunk_dataset[score["original_index"]],
+                    **{"score": score["item"]},
                 }
             )
         return retrieved_chunks

medrag_multi_modal/retrieval/medcpt_retrieval.py

@@ -231,8 +231,8 @@ class MedCPTRetriever(weave.Model):
         for score in scores:
             retrieved_chunks.append(
                 {
-                    "chunk": self._chunk_dataset[score["original_index"]],
-                    "score": score["item"],
+                    **self._chunk_dataset[score["original_index"]],
+                    **{"score": score["item"]},
                 }
             )
         return retrieved_chunks

medrag_multi_modal/retrieval/nv_embed_2.py

@@ -217,8 +217,8 @@ class NVEmbed2Retriever(weave.Model):
         for score in scores:
             retrieved_chunks.append(
                 {
-                    "chunk": self._chunk_dataset[score["original_index"]],
-                    "score": score["item"],
+                    **self._chunk_dataset[score["original_index"]],
+                    **{"score": score["item"]},
                 }
             )
         return retrieved_chunks