Update app.py

app.py

@@ -1,46 +1,60 @@
 import gradio as gr
 import json
 import random
+from datasets import load_dataset
 
-def load_jsonl(file):
-    """读取上传的jsonl文件并解析为字典列表"""
-    data = []
-    with open(file, "r", encoding="utf-8") as f:
-        for line in f:
-            data.append(json.loads(line))
-    return data
+def load_huggingface_dataset(dataset_name):
+    """从 Hugging Face 加载数据集"""
+    dataset = load_dataset(dataset_name)
+    return dataset['train']  # 假设我们使用训练集
 
-def random_data_viewer(file):
-    """读取文件并随机抽取一条数据"""
-    if file is None:
-        return "请上传一个JSONL文件！", None
+def parse_json_or_jsonl(data):
+    """解析 JSON 或 JSONL 格式的数据"""
+    parsed_data = []
+    if isinstance(data, list):  # 如果是 JSON 格式（列表）
+        parsed_data = data
+    elif isinstance(data, str):  # 如果是 JSONL 格式（每行一个 JSON 对象）
+        for line in data.splitlines():
+            if line.strip():  # 跳过空行
+                parsed_data.append(json.loads(line))
+    return parsed_data
+
+def random_data_viewer(dataset_name):
+    """从 Hugging Face 数据集中随机抽取一条数据"""
+    if dataset_name is None:
+        return "请选择一个数据集！", None
     
-    data = load_jsonl(file)
+    data = load_huggingface_dataset(dataset_name)
     if len(data) == 0:
-        return "文件为空或格式不正确！", None
+        return "数据集为空或格式不正确！", None
     
-    random_entry = random.choice(data)
-    # 格式化输出为Markdown
+    # 将数据集转换为列表形式
+    data_list = [item for item in data]
+    
+    # 随机选择一条数据
+    random_entry = random.choice(data_list)
+    
+    # 格式化输出为 Markdown
     output = "\n".join([f"**{key}**: {value}" for key, value in random_entry.items()])
-    return output, data  # 返回数据以存储供后续使用
+    return output, data_list  # 返回数据以存储供后续使用
 
 def sample_more(data):
     """从已有数据中再采样一条"""
     if not data:
-        return "没有可用数据，请先上传JSONL文件！"
+        return "没有可用数据，请先选择一个数据集！"
     
     random_entry = random.choice(data)
-    # 格式化输出为Markdown
+    # 格式化输出为 Markdown
     output = "\n".join([f"**{key}**: {value}" for key, value in random_entry.items()])
     return output
 
 # Gradio 界面
 with gr.Blocks() as app:
-    gr.Markdown("# JSONL 数据查看器")
-    gr.Markdown("上传一个JSONL文件，随机展示其中一条数据。点击按钮可以重新采样。")
+    gr.Markdown("# Hugging Face 数据集查看器")
+    gr.Markdown("选择一个 Hugging Face 数据集，随机展示其中一条数据。点击按钮可以重新采样。")
     
     with gr.Row():
-        file_upload = gr.File(file_types=[".jsonl"], label="上传JSONL文件")
+        dataset_name = gr.Textbox(label="输入 Hugging Face 数据集名称", placeholder="例如: imdb")
         sample_button = gr.Button("再采样")
     
     output_box = gr.Textbox(label="随机数据", lines=10, max_lines=20)
@@ -48,8 +62,8 @@ with gr.Blocks() as app:
     # 用于存储加载后的数据
     state_data = gr.State()
     
-    # 绑定事件：上传文件后随机取一条数据
-    file_upload.change(random_data_viewer, inputs=file_upload, outputs=[output_box, state_data])
+    # 绑定事件：输入数据集名称后随机取一条数据
+    dataset_name.change(random_data_viewer, inputs=dataset_name, outputs=[output_box, state_data])
     
     # 绑定事件：点击按钮后从已有数据中再采样
     sample_button.click(sample_more, inputs=state_data, outputs=output_box)