update dataset info

app.py

@@ -1,26 +1,60 @@
 import gradio as gr
 from arxiv_metadata_service import ArxivMetadataService
 import traceback
+import logging
+from config import DATASET_NAME
+from datasets import load_dataset
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 arxiv_service = ArxivMetadataService()
 
 def extract_metadata(query: str, max_results: int):
     try:
-        return arxiv_service.extract_and_update(query, max_results)
+        result = arxiv_service.extract_and_update(query, max_results)
+        logging.info(f"Extraction result: {result}")
+        return result
     except Exception as e:
         error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        logging.error(error_msg)
         return error_msg
 
-demo = gr.Interface(
-    fn=extract_metadata,
-    inputs=[
-        gr.Textbox(label="ArXiv Query"),
-        gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Max Results")
-    ],
-    outputs="text",
-    title="ArXiv Metadata Extractor",
-    description="Extract metadata from ArXiv papers and update the dataset."
-)
+def load_dataset_info():
+    try:
+        dataset = load_dataset(DATASET_NAME, split="train")
+        return f"Dataset contains {len(dataset)} records."
+    except Exception as e:
+        return f"Error loading dataset: {str(e)}"
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        f"""Extract metadata from ArXiv papers and update the dataset.
+        \n\nCurrently leverages the following datasets:
+        \n- [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME}/viewer) dataset.
+        """
+        )
+    
+    with gr.Tab("Extract Metadata"):
+        query_input = gr.Textbox(label="ArXiv Query")
+        max_results = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Max Results")
+        submit_button = gr.Button("Extract Metadata")
+        output = gr.Textbox(label="Result")
+        
+        submit_button.click(
+            fn=extract_metadata,
+            inputs=[query_input, max_results],
+            outputs=output
+        )
+    
+    with gr.Tab("View Dataset"):
+        refresh_button = gr.Button("Refresh Dataset Info")
+        dataset_info = gr.Textbox(label="Dataset Info")
+        
+        refresh_button.click(
+            fn=load_dataset_info,
+            inputs=[],
+            outputs=dataset_info
+        )
 
 if __name__ == "__main__":
     demo.launch()

arxiv_metadata_service.py

@@ -1,5 +1,6 @@
 from arxiv_fetcher import fetch_arxiv_metadata
 from datasets import load_dataset, Dataset
+from huggingface_hub import HfApi
 from config import DATASET_NAME
 import logging
 from typing import List, Dict, Any
@@ -7,30 +8,51 @@ from typing import List, Dict, Any
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 class ArxivMetadataService:
+    def __init__(self):
+        self.hf_api = HfApi()
+
     def extract_and_update(self, query: str, max_results: int = 10) -> str:
         metadata_list = fetch_arxiv_metadata(query, max_results)
+        if not metadata_list:
+            return "No metadata found for the given query."
         return self.update_dataset(metadata_list)
 
     def update_dataset(self, metadata_list: List[Dict[str, Any]]) -> str:
         try:
-            dataset = load_dataset(DATASET_NAME, split="train")
-            current_data = dataset.to_dict()
-            
+            # Load the existing dataset
+            try:
+                dataset = load_dataset(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 metadata_list
+            if not current_data:
+                current_data = {key: [] for key in metadata_list[0].keys()}
+
+            updated = False
             for paper in metadata_list:
-                if paper['id'] not in current_data.get('id', []):
+                entry_id = paper['entry_id'].split('/')[-1]
+                if 'entry_id' not in current_data or entry_id not in current_data['entry_id']:
+                    # Add new paper
                     for key, value in paper.items():
-                        if key not in current_data:
-                            current_data[key] = []
-                        current_data[key].append(value)
+                        current_data.setdefault(key, []).append(value)
+                    updated = True
                 else:
-                    index = current_data['id'].index(paper['id'])
+                    # Update existing paper
+                    index = current_data['entry_id'].index(entry_id)
                     for key, value in paper.items():
-                        current_data[key][index] = value
-            
-            updated_dataset = Dataset.from_dict(current_data)
-            updated_dataset.push_to_hub(DATASET_NAME, split="train")
-            
-            return f"Successfully updated dataset with {len(metadata_list)} papers"
+                        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(DATASET_NAME, split="train")
+                return f"Successfully updated dataset with {len(metadata_list)} papers"
+            else:
+                return "No new data to update."
         except Exception as e:
             logging.error(f"Failed to update dataset: {str(e)}")
             return f"Failed to update dataset: {str(e)}"

initialize_dataset.py

@@ -1,24 +1,35 @@
 from datasets import Dataset
+from huggingface_hub import HfApi
 from config import DATASET_NAME
-import huggingface_hub
+import logging
 
-# Initialize an empty dataset with the expected structure
-initial_data = {
-    "id": [],
-    "title": [],
-    "authors": [],
-    "published": [],
-    "updated": [],
-    "pdf_url": [],
-    "entry_id": [],
-    "summary": [],
-    "categories": [],
-    "primary_category": [],
-    "html_url": []
-}
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Create the dataset
-dataset = Dataset.from_dict(initial_data)
+def initialize_dataset():
+    # Initialize an empty dataset with the expected structure
+    initial_data = {
+        "entry_id": [],
+        "title": [],
+        "authors": [],
+        "published": [],
+        "updated": [],
+        "pdf_url": [],
+        "summary": [],
+        "categories": [],
+        "primary_category": [],
+        "html_url": []
+    }
 
-# Push the initial dataset to the Hub
-dataset.push_to_hub(DATASET_NAME, split="train")
+    # Create the dataset
+    dataset = Dataset.from_dict(initial_data)
+
+    try:
+        # Push the initial dataset to the Hub
+        dataset.push_to_hub(DATASET_NAME, split="train")
+        logging.info(f"Dataset {DATASET_NAME} initialized successfully with 'train' split.")
+    except Exception as e:
+        logging.error(f"Failed to initialize dataset: {str(e)}")
+        raise
+
+if __name__ == "__main__":
+    initialize_dataset()

requirements.txt

@@ -1,3 +1,5 @@
 arxiv
 datasets
 gradio
+huggingface_hub
+python-dotenv