Update app.py

app.py

@@ -993,27 +993,75 @@ def refresh_data():
 
 
 target_datasets = {
-    "aiqtech/kolaw": "https://huggingface.co/datasets/aiqtech/kolaw",
-    "heegyu/korquad-v1-v2": "https://huggingface.co/datasets/heegyu/korquad-v1-v2",
-    "nlpai-lab/kowiki": "https://huggingface.co/datasets/nlpai-lab/kowiki",
-    "KETI-AIR/korpora": "https://huggingface.co/datasets/KETI-AIR/korpora",
-    "heegyu/korean-parallel-corpora": "https://huggingface.co/datasets/heegyu/korean-parallel-corpora",
-    "heegyu/korean-hate-speech": "https://huggingface.co/datasets/heegyu/korean-hate-speech",
-    "KETI-AIR/korean-parallel-corpora": "https://huggingface.co/datasets/KETI-AIR/korean-parallel-corpora",
-    "heegyu/korean-chatbot-data": "https://huggingface.co/datasets/heegyu/korean-chatbot-data",
-    "heegyu/korean-qa": "https://huggingface.co/datasets/heegyu/korean-qa",
-    "heegyu/korean-summarization": "https://huggingface.co/datasets/heegyu/korean-summarization",
-    "nlpai-lab/kullm-chat-v2": "https://huggingface.co/datasets/nlpai-lab/kullm-chat-v2",
-    "upstage/open-ko-llm-leaderboard": "https://huggingface.co/datasets/upstage/open-ko-llm-leaderboard"
+    "aiqtech/kolaw": "https://huggingface.co/datasets/aiqtech/kolaw"
+    # 필요한 데이터셋 추가
 }
 
+def get_korea_datasets():
+    """Korea 관련 데이터셋 검색"""
+    params = {
+        "search": "korea",
+        "full": "True",
+        "limit": 1000
+    }
+    
+    try:
+        response = requests.get(
+            "https://huggingface.co/api/datasets",
+            headers={'Authorization': f'Bearer {HF_TOKEN}'},
+            params=params
+        )
+        
+        if response.status_code == 200:
+            return response.json()
+        else:
+            print(f"Failed to fetch Korea datasets: {response.status_code}")
+            return []
+    except Exception as e:
+        print(f"Error fetching Korea datasets: {str(e)}")
+        return []
+
+def get_all_datasets(limit=3000):
+    """모든 데이터셋과 Korea 관련 데이터셋 가져오기"""
+    all_datasets = []
+    page_size = 1000
+    
+    for offset in range(0, limit, page_size):
+        params = {
+            'limit': min(page_size, limit - offset),
+            'full': 'True',
+            'offset': offset
+        }
+        
+        response = requests.get(
+            "https://huggingface.co/api/datasets",
+            headers={'Authorization': f'Bearer {HF_TOKEN}'},
+            params=params
+        )
+        
+        if response.status_code == 200:
+            all_datasets.extend(response.json())
+            print(f"Fetched datasets {offset+1} to {offset+len(response.json())}")
+        else:
+            print(f"Failed to fetch datasets at offset {offset}: {response.status_code}")
+            break
+    
+    # Korea 검색 결과 추가
+    korea_datasets = get_korea_datasets()
+    existing_ids = {dataset.get('id', '') for dataset in all_datasets}
+    
+    for korea_dataset in korea_datasets:
+        if korea_dataset.get('id', '') not in existing_ids:
+            all_datasets.append(korea_dataset)
+            existing_ids.add(korea_dataset.get('id', ''))
+    
+    return all_datasets[:limit]
+
 def get_datasets_data(progress=gr.Progress()):
     def calculate_rank(dataset_id, all_global_datasets, korea_datasets):
-        # 글로벌 순위 확인
         global_rank = next((idx for idx, d in enumerate(all_global_datasets, 1) 
                           if d.get('id', '').strip() == dataset_id.strip()), None)
         
-        # Korea 데이터셋인 경우
         is_korea = any(d.get('id', '').strip() == dataset_id.strip() for d in korea_datasets)
         
         if is_korea:
@@ -1039,37 +1087,9 @@ def get_datasets_data(progress=gr.Progress()):
             empty_df = pd.DataFrame(columns=['Global Rank', 'Dataset ID', 'Title', 'Downloads', 'Likes', 'Korea Search', 'URL'])
             return fig, error_html, empty_df
 
-        # 일반 데이터셋과 Korea 관련 데이터셋 가져오기
-        params = {
-            'limit': 3000,
-            'full': 'true'
-        }
-        
-        all_datasets_response = requests.get(
-            "https://huggingface.co/api/datasets",
-            headers={'Authorization': f'Bearer {HF_TOKEN}'},
-            params=params
-        )
-        
-        korea_params = {
-            'search': 'korea',
-            'limit': 3000,
-            'full': 'true'
-        }
+        all_global_datasets = get_all_datasets(limit=3000)
+        korea_datasets = get_korea_datasets()
         
-        korea_datasets_response = requests.get(
-            "https://huggingface.co/api/datasets",
-            headers={'Authorization': f'Bearer {HF_TOKEN}'},
-            params=korea_params
-        )
-        
-        all_global_datasets = all_datasets_response.json()
-        korea_datasets = korea_datasets_response.json()
-        
-        # 시각화를 위한 Figure 생성
-        fig = go.Figure()
-        
-        # 순위 정보 수집
         filtered_datasets = []
         for dataset_id in target_datasets.keys():
             try:
@@ -1092,8 +1112,6 @@ def get_datasets_data(progress=gr.Progress()):
                         'title': dataset_data.get('title', 'No Title'),
                         'is_korea': is_korea
                     })
-                    
-                    print(f"Dataset {dataset_id}: Rank={rank}, Is Korea={is_korea}")
                 else:
                     filtered_datasets.append({
                         'id': dataset_id,
@@ -1107,12 +1125,13 @@ def get_datasets_data(progress=gr.Progress()):
                 print(f"Error processing {dataset_id}: {str(e)}")
                 continue
 
-        # 순위로 정렬
         filtered_datasets.sort(key=lambda x: float('inf') if isinstance(x['global_rank'], str) else x['global_rank'])
         
         # 시각화 데이터 준비
         valid_datasets = [d for d in filtered_datasets if isinstance(d['global_rank'], (int, float))]
         
+        fig = go.Figure()
+        
         if valid_datasets:
             ids = [d['id'] for d in valid_datasets]
             ranks = [d['global_rank'] for d in valid_datasets]
@@ -1120,10 +1139,12 @@ def get_datasets_data(progress=gr.Progress()):
             fig.add_trace(go.Bar(
                 x=ids,
                 y=[3001 - r for r in ranks],
-                text=[f"Rank: #{r}<br>Downloads: {format(d['downloads'], ',')}<br>Likes: {format(d['likes'], ',')}"
+                text=[f"Rank: #{r}<br>{'🇰🇷 Korea Dataset<br>' if d['is_korea'] else ''}"
+                      f"Downloads: {format(d['downloads'], ',')}<br>"
+                      f"Likes: {format(d['likes'], ',')}"
                       for r, d in zip(ranks, valid_datasets)],
                 textposition='auto',
-                marker_color=['rgba(255,0,0,0.6)' if d['is_korea'] else 'rgba(0,0,255,0.6)' 
+                marker_color=['rgba(255,0,0,0.6)' if d['is_korea'] else 'rgba(0,0,255,0.6)'
                             for d in valid_datasets],
                 opacity=0.8
             ))
@@ -1198,8 +1219,6 @@ def get_datasets_data(progress=gr.Progress()):
         } for d in filtered_datasets])
         
         progress(1.0, desc="Complete!")
-
-
         return fig, html_content, df
 
     except Exception as e: