modify app.py

app.py

@@ -39,23 +39,21 @@ class MGTHuman(datasets.GeneratorBasedBuilder):
 
     def get_text_by_index(self, filepath, index):
         count = 0
-        for file in filepath:
-            with open(file, 'r') as f:
-                data = json.load(f)
-                
-            for row in data:
-                if not row["text"].strip():
-                    continue
-                if count == index:
-                    text = self.truncate_text(row["text"], max_tokens=2048)
-                    return text
-                count += 1
+        with open(filepath, 'r') as f:
+            data = json.load(f)
+        for row in data:
+            if not row["text"].strip():
+                continue
+            if count == index:
+                text = self.truncate_text(row["text"], max_tokens=2048)
+                return text
+            count += 1
         return "Index 超出范围，请输入有效的数字。"
 
 # Streamlit UI
 st.title("MGTHuman Dataset Viewer")
 
-# 文件夹上传
+# 上传包含 JSON 文件的 ZIP 文件
 uploaded_folder = st.file_uploader("上传包含 JSON 文件的 ZIP 文件夹", type=["zip"])
 if uploaded_folder:
     folder_path = Path("temp")
@@ -67,17 +65,26 @@ if uploaded_folder:
     with zipfile.ZipFile(zip_path, 'r') as zip_ref:
         zip_ref.extractall(folder_path)
 
-    # 获取解压后的所有 JSON 文件路径
-    json_files = list(folder_path.glob("*.json"))
-    
-    # 选择数据配置
-    config_name = st.selectbox("选择数据配置", ["human", "Moonshot", "gpt35", "Llama3", "Mixtral", "Qwen"])
-    mgt_human = MGTHuman(name=config_name)
+    # 获取所有 JSON 文件并分类到不同的 domain
+    category = {}
+    for json_file in folder_path.glob("*.json"):
+        domain = json_file.stem.split('_task3')[0]
+        category.setdefault(domain, []).append(str(json_file))
+
+    # 显示可用的 domain
+    st.write("可用的数据种类:", list(category.keys()))
 
+    # 用户选择 domain
+    selected_domain = st.selectbox("选择数据种类", options=list(category.keys()))
+    
     # 输入序号查看文本
     index_to_view = st.number_input("输入要查看的文本序号", min_value=0, step=1)
+    
     if st.button("显示文本"):
-        text = mgt_human.get_text_by_index(json_files, index=index_to_view)
+        # 选择第一个文件进行展示
+        file_to_display = category[selected_domain][0]
+        mgt_human = MGTHuman(name=selected_domain)
+        text = mgt_human.get_text_by_index(file_to_display, index=index_to_view)
         st.write("对应的文本内容为：", text)
 
 # 清理上传文件的临时目录