geekyrakshit commited on
Commit
1f48fed
·
1 Parent(s): 2b64a07

update: retrieved chunks data structure

Browse files
medrag_multi_modal/retrieval/bm25s_retrieval.py CHANGED
@@ -175,7 +175,7 @@ class BM25sRetriever(weave.Model):
175
  results.documents.flatten().tolist(),
176
  results.scores.flatten().tolist(),
177
  ):
178
- retrieved_chunks.append({"chunk": chunk, "score": score})
179
  return retrieved_chunks
180
 
181
  @weave.op()
 
175
  results.documents.flatten().tolist(),
176
  results.scores.flatten().tolist(),
177
  ):
178
+ retrieved_chunks.append({**chunk, **{"score": score}})
179
  return retrieved_chunks
180
 
181
  @weave.op()
medrag_multi_modal/retrieval/contriever_retrieval.py CHANGED
@@ -192,8 +192,8 @@ class ContrieverRetriever(weave.Model):
192
  for score in scores:
193
  retrieved_chunks.append(
194
  {
195
- "chunk": self._chunk_dataset[score["original_index"]],
196
- "score": score["item"],
197
  }
198
  )
199
  return retrieved_chunks
 
192
  for score in scores:
193
  retrieved_chunks.append(
194
  {
195
+ **self._chunk_dataset[score["original_index"]],
196
+ **{"score": score["item"]},
197
  }
198
  )
199
  return retrieved_chunks
medrag_multi_modal/retrieval/medcpt_retrieval.py CHANGED
@@ -231,8 +231,8 @@ class MedCPTRetriever(weave.Model):
231
  for score in scores:
232
  retrieved_chunks.append(
233
  {
234
- "chunk": self._chunk_dataset[score["original_index"]],
235
- "score": score["item"],
236
  }
237
  )
238
  return retrieved_chunks
 
231
  for score in scores:
232
  retrieved_chunks.append(
233
  {
234
+ **self._chunk_dataset[score["original_index"]],
235
+ **{"score": score["item"]},
236
  }
237
  )
238
  return retrieved_chunks
medrag_multi_modal/retrieval/nv_embed_2.py CHANGED
@@ -217,8 +217,8 @@ class NVEmbed2Retriever(weave.Model):
217
  for score in scores:
218
  retrieved_chunks.append(
219
  {
220
- "chunk": self._chunk_dataset[score["original_index"]],
221
- "score": score["item"],
222
  }
223
  )
224
  return retrieved_chunks
 
217
  for score in scores:
218
  retrieved_chunks.append(
219
  {
220
+ **self._chunk_dataset[score["original_index"]],
221
+ **{"score": score["item"]},
222
  }
223
  )
224
  return retrieved_chunks