Spaces:
Paused
Paused
File size: 1,975 Bytes
97e8d87 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from typing import List, Dict, Any
from datasets import load_dataset, Dataset
class DatasetManagementService:
def __init__(self, dataset_name: str):
self.dataset_name = dataset_name
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()
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']:
for key, value in paper.items():
current_data.setdefault(key, []).append(value)
updated = True
else:
index = current_data['entry_id'].index(entry_id)
for key, value in paper.items():
if current_data[key][index] != value:
current_data[key][index] = value
updated = True
if updated:
updated_dataset = Dataset.from_dict(current_data)
updated_dataset.push_to_hub(self.dataset_name, split="train")
return f"Successfully updated dataset with {len(new_metadata)} papers"
else:
return "No new data to update."
except Exception as e:
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")
# Usage:
# dataset_service = DatasetManagementService("your_dataset_name")
# result = dataset_service.update_dataset(new_metadata)
# records = dataset_service.get_dataset_records()
|