adding components

app.py

@@ -1,133 +1,38 @@
 import gradio as gr
-import arxiv
-import traceback
-import logging
 from typing import List, Dict, Any
-from datasets import load_dataset, Dataset
-from huggingface_hub import HfApi
 from config import DATASET_NAME
+from arxiv_retrieval_service import ArxivRetrievalService
+from dataset_management_service import DatasetManagementService
 
-# Logging setup
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Initialize services
+arxiv_service = ArxivRetrievalService()
+dataset_service = DatasetManagementService(DATASET_NAME)
 
-# Arxiv Fetcher logic
-def fetch_metadata(query: str, max_results: int = 10) -> List[Dict[str, Any]]:
-    logging.info(f"Fetching arXiv metadata for query: {query}")
-    if not query.strip():
-        logging.warning("Empty or whitespace-only query provided")
-        return []
-    
-    client = arxiv.Client(page_size=max_results, delay_seconds=3, num_retries=3)
-    search = arxiv.Search(query=query, max_results=max_results, sort_by=arxiv.SortCriterion.SubmittedDate)
-    
-    results = []
+def handle_metadata_extraction(query: str, max_results: int) -> str:
     try:
-        for result in client.results(search):
-            metadata = {
-                "title": result.title,
-                "authors": [author.name for author in result.authors],
-                "published": result.published.isoformat(),
-                "updated": result.updated.isoformat(),
-                "pdf_url": result.pdf_url,
-                "entry_id": result.entry_id,
-                "summary": result.summary,
-                "categories": result.categories,
-                "primary_category": result.primary_category,
-                "html_url": f"http://arxiv.org/abs/{result.entry_id.split('/')[-1]}"
-            }
-            results.append(metadata)
-        logging.info(f"Fetched metadata for {len(results)} papers")
-    except Exception as e:
-        logging.error(f"Error fetching metadata: {str(e)}")
-    
-    return results
-
-# Arxiv Metadata Service logic
-class ArxivMetadataService:
-    def __init__(self):
-        self.hf_api = HfApi()
-
-    def extract_metadata_and_update_dataset(self, query: str, max_results: int = 10) -> str:
-        metadata_list = fetch_metadata(query, max_results)
+        # Fetch metadata from ArXiv
+        metadata_list = arxiv_service.fetch_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:
-            # 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:
-                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():
-                        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:
-                            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)}"
-
-    def get_dataset_records(self):
-        try:
-            dataset = load_dataset(DATASET_NAME, split="train")
-            records = dataset.to_pandas().to_dict(orient="records")
-            return records
-        except Exception as e:
-            return f"Error loading dataset: {str(e)}"
-
-# Initialize Arxiv Metadata Service
-arxiv_service = ArxivMetadataService()
-
-# Define Gradio functions
-def handle_metadata_extraction(query: str, max_results: int):
-    try:
-        result = arxiv_service.extract_metadata_and_update_dataset(query, max_results)
-        logging.info(f"Extraction result: {result}")
+        
+        # Update the dataset with new metadata
+        result = dataset_service.update_dataset(metadata_list)
         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
+        return f"An error occurred: {str(e)}"
 
-def handle_dataset_view():
+def handle_dataset_view() -> List[Dict[str, Any]]:
     try:
-        records = arxiv_service.get_dataset_records()
-        return records
+        return dataset_service.get_dataset_records()
     except Exception as e:
-        return f"Error loading dataset: {str(e)}"
+        return [{"error": f"Error loading dataset: {str(e)}"}]
 
 # 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 datasets:
-        \n- [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME}/viewer) dataset.
+        \n\nCurrently leverages the following dataset:
+        \n- [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME}/viewer)
         """
     )
     
@@ -154,5 +59,4 @@ with gr.Blocks() as demo:
         )
 
 if __name__ == "__main__":
-    demo.queue()
-    demo.launch()
+    demo.launch()

arxiv_retrieval_service.py

@@ -0,0 +1,35 @@
+import arxiv
+from typing import List, Dict, Any
+
+class ArxivRetrievalService:
+    def __init__(self):
+        self.client = arxiv.Client(delay_seconds=3, num_retries=3)
+
+    def fetch_metadata(self, query: str, max_results: int = 10) -> List[Dict[str, Any]]:
+        search = arxiv.Search(
+            query=query, 
+            max_results=max_results, 
+            sort_by=arxiv.SortCriterion.SubmittedDate
+        )
+        
+        results = []
+        for result in self.client.results(search):
+            metadata = {
+                "title": result.title,
+                "authors": [author.name for author in result.authors],
+                "published": result.published.isoformat(),
+                "updated": result.updated.isoformat(),
+                "pdf_url": result.pdf_url,
+                "entry_id": result.entry_id,
+                "summary": result.summary,
+                "categories": result.categories,
+                "primary_category": result.primary_category,
+                "html_url": f"http://arxiv.org/abs/{result.entry_id.split('/')[-1]}"
+            }
+            results.append(metadata)
+        
+        return results
+
+# Usage:
+# arxiv_service = ArxivRetrievalService()
+# metadata = arxiv_service.fetch_metadata("quantum computing", max_results=5)

dataset_management_service.py

@@ -0,0 +1,46 @@
+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()