Update app.py

app.py

@@ -13,7 +13,6 @@ import datasets
 import requests
 
 from datetime import timezone  # Ensure timezone is imported
-import html  # For sanitizing user inputs
 
 # --- Data Loading and Processing ---
 
@@ -38,10 +37,6 @@ def get_df() -> pd.DataFrame:
     df["date"] = pd.to_datetime(df["date"], errors='coerce')
     df["date"] = df["date"].dt.strftime("%Y-%m-%d").fillna(datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d"))
     
-    # Ensure 'upvotes' and 'num_comments' are present and fill NaNs with 0
-    df['upvotes'] = df.get('upvotes', 0).fillna(0)
-    df['num_comments'] = df.get('num_comments', 0).fillna(0)
-    
     # Prepare the DataFrame by removing 'abstract'
     paper_info = []
     for _, row in tqdm.auto.tqdm(df.iterrows(), total=len(df)):
@@ -61,18 +56,8 @@ class Prettifier:
     """
     Converts raw DataFrame rows into a prettified format suitable for display.
     """
-    def __init__(self):
-        self.expected_columns = [
-            "arxiv_id",
-            "date_display",
-            "date",
-            "paper_page",
-            "title",
-            "github",
-            "👍",
-            "💬",
-        ]
-
+    REQUIRED_COLUMNS = ["arxiv_id", "date_display", "date", "paper_page", "title", "github", "👍", "💬"]
+    
     @staticmethod
     def get_github_link(link: str) -> str:
         if not link:
@@ -87,23 +72,25 @@ class Prettifier:
         new_rows = []
         for _, row in df.iterrows():
             # Handle date_display as a clickable link
-            date_display = self.create_link(row.date, f"https://huggingface.co/papers?date={row.date}")
+            date_display = Prettifier.create_link(row.get("date", ""), f"https://huggingface.co/papers?date={row.get('date', '')}")
             
             new_row = {
-                "arxiv_id": row.get("arxiv_id", ""),                        # Include arxiv_id
-                "date_display": date_display,                               # For display
+                "arxiv_id": row.get("arxiv_id", ""),  # Include arxiv_id
+                "date_display": date_display,         # For display
                 "date": row.get("date", datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d")),  # For internal calculations
-                "paper_page": self.create_link(row.get("arxiv_id", ""), row.get("paper_page", "#")),
+                "paper_page": Prettifier.create_link(row.get("arxiv_id", ""), row.get("paper_page", "#")),
                 "title": row.get("title", "No title"),
-                "github": self.get_github_link(row.get("github", "")),
+                "github": Prettifier.get_github_link(row.get("github", "")),
                 "👍": row.get("upvotes", 0),
                 "💬": row.get("num_comments", 0),
             }
             new_rows.append(new_row)
+        
+        # If no rows, return empty DataFrame with required columns to prevent KeyError
         if not new_rows:
-            return pd.DataFrame(columns=self.expected_columns)
-        else:
-            return pd.DataFrame(new_rows)
+            return pd.DataFrame(columns=self.REQUIRED_COLUMNS)
+        
+        return pd.DataFrame(new_rows)
 
 
 class PaperList:
@@ -128,11 +115,11 @@ class PaperList:
 
     @property
     def column_names(self):
-        return list(map(operator.itemgetter(0), self.COLUMN_INFO))
+        return [col[0] for col in self.COLUMN_INFO]
 
     @property
     def column_datatype(self):
-        return list(map(operator.itemgetter(1), self.COLUMN_INFO))
+        return [col[1] for col in self.COLUMN_INFO]
 
     def search(
         self,
@@ -225,11 +212,9 @@ class PaperManager:
             
             # Convert 'date' column to datetime
             df_sorted = df.copy()
-            df_sorted['date_parsed'] = pd.to_datetime(df_sorted['date'], errors='coerce').dt.tz_localize(timezone.utc, ambiguous='NaT', nonexistent='NaT')
-            # Handle NaT values by setting them to a very old date
-            df_sorted['date_parsed'] = df_sorted['date_parsed'].fillna(datetime.datetime.min.replace(tzinfo=timezone.utc))
+            df_sorted['date_parsed'] = pd.to_datetime(df_sorted['date'], errors='coerce').dt.tz_localize(timezone.utc)
             df_sorted = df_sorted[df_sorted['date_parsed'] >= time_threshold]
-            df_sorted = df_sorted.sort_values(by='👍', ascending=False).drop(columns=['date_parsed'])
+            df_sorted = df_sorted.sort_values(by='upvotes', ascending=False).drop(columns=['date_parsed'])
         else:
             df_sorted = df
 
@@ -257,9 +242,8 @@ class PaperManager:
         """
         Sets the current search query and re-sorts the papers.
         """
-        sanitized_query = html.escape(query)  # Sanitize user input
-        print(f"Setting search query to: {sanitized_query}")
-        self.current_search_query = sanitized_query
+        print(f"Setting search query to: {query}")
+        self.current_search_query = query
         self.sort_papers()
         return True  # Assume success
 
@@ -397,7 +381,10 @@ def change_sort_method_ui(method: str, time_frame: str = "all time") -> str:
     """
     Changes the sort method and, if 'top' is selected, sets the time frame.
     """
-    paper_manager.set_sort_method(method.lower(), time_frame if method.lower() == "top" else None)
+    if method.lower() == "top":
+        paper_manager.set_sort_method(method.lower(), time_frame)
+    else:
+        paper_manager.set_sort_method(method.lower())
     return paper_manager.get_current_page_papers()
 
 
@@ -702,4 +689,4 @@ with demo:
 # --- Launch the App ---
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()