import streamlit as st
import pandas as pd
from datasets import load_dataset

# 1. 设置页面标题
st.set_page_config(page_title="我的数据处理工具", layout="wide")
st.title("📊 数据集预览与处理工具")

# 2. 加载数据 (以 arena-human-preference-140k 为例，为了速度只加载前100行)
@st.cache_data # 这个装饰器能让数据加载后缓存，避免每次刷新页面都重新下载
def load_hf_data():
    dataset = load_dataset("lmarena-ai/arena-human-preference-140k", split="train", streaming=True)
    # 取前100条转为 Pandas 表格，方便展示
    df = pd.DataFrame(list(dataset.take(100)))
    return df

df = load_hf_data()

# 3. 侧边栏：选择行数
st.sidebar.header("数据筛选")
row_index = st.sidebar.number_input("选择行号", min_value=0, max_value=len(df)-1, value=0)

# 4. 主界面：展示表格
st.subheader("数据集概览")
st.dataframe(df.head(10)) # 展示前10行

# 5. 详细解析特定的一行
st.divider()
st.subheader(f"第 {row_index} 行详细内容")

selected_row = df.iloc[row_index]

col1, col2 = st.columns(2)
with col1:
    st.info("模型 A 的回答")
    st.write(selected_row['model_a'])

with col2:
    st.info("模型 B 的回答")
    st.write(selected_row['model_b'])

# 6. 数据处理按钮
if st.button("🚀 标记为高质量数据"):
    st.success(f"已记录第 {row_index} 行！(此处可扩展保存逻辑)")