simplify app.py

app.py

@@ -1,17 +1,113 @@
 import gradio as gr
-from arxiv_metadata_service import ArxivMetadataService
+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 datasets import load_dataset
 
+# Logging setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
+# 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 = []
+    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)
+        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()
 
-def extract_metadata(query: str, max_results: int):
+# Define Gradio functions
+def handle_metadata_extraction(query: str, max_results: int):
     try:
-        result = arxiv_service.extract_and_update(query, max_results)
+        result = arxiv_service.extract_metadata_and_update_dataset(query, max_results)
         logging.info(f"Extraction result: {result}")
         return result
     except Exception as e:
@@ -19,20 +115,21 @@ def extract_metadata(query: str, max_results: int):
         logging.error(error_msg)
         return error_msg
 
-def load_dataset_info():
+def handle_dataset_view():
     try:
-        dataset = load_dataset(DATASET_NAME, split="train")
-        return f"Dataset contains {len(dataset)} records."
+        records = arxiv_service.get_dataset_records()
+        return records
     except Exception as e:
         return 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.
         """
-        )
+    )
     
     with gr.Tab("Extract Metadata"):
         query_input = gr.Textbox(label="ArXiv Query")
@@ -41,20 +138,21 @@ with gr.Blocks() as demo:
         output = gr.Textbox(label="Result")
         
         submit_button.click(
-            fn=extract_metadata,
+            fn=handle_metadata_extraction,
             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")
+        dataset_info = gr.JSON(label="Dataset Info")
         
         refresh_button.click(
-            fn=load_dataset_info,
+            fn=handle_dataset_view,
             inputs=[],
             outputs=dataset_info
         )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue()
+    demo.launch()