update ds mgmt

dataset_management_service.py

@@ -7,20 +7,31 @@ class DatasetManagementService:
 
     def update_dataset(self, new_metadata: List[Dict[str, Any]]) -> str:
         try:
-            dataset = load_dataset(self.dataset_name, split="train")
-            current_data = dataset.to_dict()
+            # Try to load the existing dataset
+            try:
+                dataset = load_dataset(self.dataset_name, split="train")
+                current_data = dataset.to_dict()
+            except Exception:
+                # If loading fails, start with an empty dictionary
+                current_data = {}
 
+            # If the dataset is empty, initialize it with the structure from new_metadata
             if not current_data:
                 current_data = {key: [] for key in new_metadata[0].keys()}
 
             updated = False
             for paper in new_metadata:
-                entry_id = paper['entry_id'].split('/')[-1]
-                if 'entry_id' not in current_data or entry_id not in current_data['entry_id']:
+                entry_id = paper['entry_id']
+                if 'entry_id' not in current_data:
+                    current_data['entry_id'] = []
+
+                if entry_id not in current_data['entry_id']:
+                    # Add new paper
                     for key, value in paper.items():
                         current_data.setdefault(key, []).append(value)
                     updated = True
                 else:
+                    # Update existing paper
                     index = current_data['entry_id'].index(entry_id)
                     for key, value in paper.items():
                         if current_data[key][index] != value:
@@ -37,10 +48,15 @@ class DatasetManagementService:
             return f"Failed to update dataset: {str(e)}"
 
     def get_dataset_records(self) -> List[Dict[str, Any]]:
-        dataset = load_dataset(self.dataset_name, split="train")
-        return dataset.to_pandas().to_dict(orient="records")
+        try:
+            dataset = load_dataset(self.dataset_name, split="train")
+            if len(dataset) == 0:
+                return []
+            return dataset.to_pandas().to_dict(orient="records")
+        except Exception as e:
+            return [{"error": f"Error loading dataset: {str(e)}"}]
 
 # Usage:
 # dataset_service = DatasetManagementService("your_dataset_name")
 # result = dataset_service.update_dataset(new_metadata)
-# records = dataset_service.get_dataset_records()
+# records = dataset_service.get_dataset_records()