From dbc96d491fef4dbe6eca1812abde37b94ccd679a Mon Sep 17 00:00:00 2001 From: Chen Qian Date: Fri, 20 Dec 2024 13:16:34 -0800 Subject: [PATCH] fix databricksrm (#1967) --- dspy/retrieve/databricks_rm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dspy/retrieve/databricks_rm.py b/dspy/retrieve/databricks_rm.py index d9eb45210..249b3856d 100644 --- a/dspy/retrieve/databricks_rm.py +++ b/dspy/retrieve/databricks_rm.py @@ -288,7 +288,7 @@ def forward( page_content=doc[self.text_column_name], metadata={ "doc_id": self._extract_doc_ids(doc), - "doc_uri": doc[self.docs_uri_column_name], + "doc_uri": doc[self.docs_uri_column_name] if self.docs_uri_column_name else None, } | self._get_extra_columns(doc), type="Document", @@ -300,7 +300,7 @@ def forward( return Prediction( docs=[doc[self.text_column_name] for doc in sorted_docs], doc_ids=[self._extract_doc_ids(doc) for doc in sorted_docs], - doc_uris=[doc[self.docs_uri_column_name] for doc in sorted_docs], + doc_uris=[doc[self.docs_uri_column_name] for doc in sorted_docs] if self.docs_uri_column_name else None, extra_columns=[self._get_extra_columns(item) for item in sorted_docs], )