import gradio as gr
import json
import random
from datasets import load_dataset

def load_huggingface_dataset(dataset_name):
    """从 Hugging Face 加载数据集"""
    dataset = load_dataset(dataset_name)
    return dataset['train']  # 假设我们使用训练集

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_huggingface_dataset(dataset_name)
    if len(data) == 0:
        return "数据集为空或格式不正确！", None
    
    # 将数据集转换为列表形式
    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_list  # 返回数据以存储供后续使用

def sample_more(data):
    """从已有数据中再采样一条"""
    if not data:
        return "没有可用数据，请先选择一个数据集！"
    
    random_entry = random.choice(data)
    # 格式化输出为 Markdown
    output = "\n".join([f"**{key}**: {value}" for key, value in random_entry.items()])
    return output

# Gradio 界面
with gr.Blocks() as app:
    gr.Markdown("# Hugging Face 数据集查看器")
    gr.Markdown("选择一个 Hugging Face 数据集，随机展示其中一条数据。点击按钮可以重新采样。")
    
    with gr.Row():
        dataset_name = gr.Textbox(label="输入 Hugging Face 数据集名称", placeholder="例如: imdb")
        sample_button = gr.Button("再采样")
    
    output_box = gr.Textbox(label="随机数据", lines=10, max_lines=20)
    
    # 用于存储加载后的数据
    state_data = gr.State()
    
    # 绑定事件：输入数据集名称后随机取一条数据
    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)

# 启动 Gradio 应用
app.launch()