fix return datatype in _sample_with_metadata (#145)

hitsz-ids · Feb 27, 2024 · ad3d188 · ad3d188
1 parent f1241b5
commit ad3d188
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/sdgx/models/LLM/single_table/gpt.py b/sdgx/models/LLM/single_table/gpt.py
@@ -468,7 +468,11 @@ def _sample_with_metadata(self, count, *args, **kwargs):
             # Update the remaining count
             remaining_cnt = remaining_cnt - current_cnt
 
-        return count  # Return the input count
+        self._result_list.append(result)
+
+        # return result in pd.DataFrame
+        final_columns = self.columns + self.off_table_features
+        return pd.DataFrame(self._result_list, columns=final_columns)
 
     def _sample_with_data(self, count, *args, **kwargs):
         """