update get_dataset_records

app.py

@@ -3,36 +3,56 @@ from typing import List, Dict, Any
 from config import DATASET_NAME
 from arxiv_retrieval_service import ArxivRetrievalService
 from dataset_management_service import DatasetManagementService
+import logging
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Initialize services
 arxiv_service = ArxivRetrievalService()
 dataset_service = DatasetManagementService(DATASET_NAME)
 
 def handle_metadata_extraction(query: str, max_results: int) -> str:
     try:
-        # Fetch metadata from ArXiv
+        logging.info(f"Fetching metadata for query: {query}, max_results: {max_results}")
         metadata_list = arxiv_service.fetch_metadata(query, max_results)
         if not metadata_list:
             return "No metadata found for the given query."
         
-        # Update the dataset with new metadata
         result = dataset_service.update_dataset(metadata_list)
+        logging.info(f"Dataset update result: {result}")
         return result
     except Exception as e:
-        return f"An error occurred: {str(e)}"
+        error_msg = f"An error occurred during metadata extraction: {str(e)}"
+        logging.error(error_msg)
+        return error_msg
 
-def handle_dataset_view() -> List[Dict[str, Any]]:
+def handle_dataset_view(page: int = 1, page_size: int = 10) -> Dict[str, Any]:
+    logging.info(f"handle_dataset_view called with page={page}, page_size={page_size}")
     try:
-        return dataset_service.get_dataset_records()
+        total_records = dataset_service.get_dataset_size()
+        logging.info(f"Total records: {total_records}")
+        
+        records = dataset_service.get_dataset_records(page, page_size)
+        logging.info(f"Records type: {type(records)}")
+        logging.info(f"Number of records returned: {len(records)}")
+        
+        result = {
+            "total_records": total_records,
+            "current_page": page,
+            "records": records
+        }
+        logging.info(f"Returning result: {result}")
+        return result
     except Exception as e:
-        return [{"error": f"Error loading dataset: {str(e)}"}]
+        error_msg = f"Error loading dataset: {str(e)}"
+        logging.error(error_msg)
+        return {"error": error_msg}
 
-# Define Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown(
-        f"""Extract metadata from ArXiv papers and update the dataset.
-        \n\nCurrently leverages the following dataset:
-        \n- [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME}/viewer)
+        f"""# ArXiv Metadata Extraction and Dataset Management
+        
+        This application extracts metadata from ArXiv papers and manages the dataset:
+        [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME}/viewer)
         """
     )
     
@@ -49,12 +69,13 @@ with gr.Blocks() as demo:
         )
     
     with gr.Tab("View Dataset"):
-        refresh_button = gr.Button("Refresh Dataset Info")
+        page_number = gr.Number(value=1, label="Page Number", precision=0)
+        refresh_button = gr.Button("Refresh Dataset View")
         dataset_info = gr.JSON(label="Dataset Info")
         
         refresh_button.click(
             fn=handle_dataset_view,
-            inputs=[],
+            inputs=[page_number],
             outputs=dataset_info
         )
 

dataset_management_service.py

@@ -1,5 +1,6 @@
 from typing import List, Dict, Any
 from datasets import load_dataset, Dataset
+import logging
 
 class DatasetManagementService:
     def __init__(self, dataset_name: str):
@@ -47,13 +48,31 @@ class DatasetManagementService:
         except Exception as e:
             return f"Failed to update dataset: {str(e)}"
 
-    def get_dataset_records(self) -> List[Dict[str, Any]]:
+    def get_dataset_size(self) -> int:
         try:
             dataset = load_dataset(self.dataset_name, split="train")
-            if len(dataset) == 0:
-                return []
-            return dataset.to_pandas().to_dict(orient="records")
+            size = len(dataset)
+            logging.info(f"Dataset size: {size}")
+            return size
         except Exception as e:
+            logging.error(f"Error getting dataset size: {str(e)}")
+            return 0
+
+    def get_dataset_records(self, page: int, page_size: int) -> List[Dict[str, Any]]:
+        try:
+            dataset = load_dataset(self.dataset_name, split="train")
+            start_idx = (page - 1) * page_size
+            end_idx = start_idx + page_size
+            records = dataset[start_idx:end_idx]
+            
+            # Convert to list of dictionaries
+            records_list = [dict(zip(records.keys(), values)) for values in zip(*records.values())]
+            
+            logging.info(f"Records type: {type(records_list)}")
+            logging.info(f"Number of records: {len(records_list)}")
+            return records_list
+        except Exception as e:
+            logging.error(f"Error loading dataset records: {str(e)}")
             return [{"error": f"Error loading dataset: {str(e)}"}]
 
 # Usage: