Update app.py

app.py

@@ -3,76 +3,52 @@ import zipfile
 import os
 import shutil
 from datasets import load_dataset
+import uuid
 
 def zip_folder(dataset_name, folder_name, output_zip):
-    """
-    指定されたHugging Faceデータセット内のフォルダをZIPに圧縮する。
-
-    Args:
-        dataset_name (str): データセット名。
-        folder_name (str): 圧縮するフォルダの名前。
-        output_zip (str): 出力するZIPファイル名。
-
-    Returns:
-        str: ZIPファイルのパス。
-    """
     try:
-        # データセットをロード
         print(f"ロード対象データセット: {dataset_name}")
         dataset = load_dataset(dataset_name)
-    except FileNotFoundError:
-        return f"指定されたデータセット '{dataset_name}' が見つかりません。"
     except Exception as e:
-        return f"エラーが発生しました: {str(e)}"
+        return f"データセットをロードできませんでした: {str(e)}"
 
-    # 一時ディレクトリに保存
-    temp_dir = "temp_dataset"
-    if os.path.exists(temp_dir):
-        shutil.rmtree(temp_dir)
-    os.makedirs(temp_dir)
+    if not hasattr(dataset, "save_to_disk"):
+        return f"データセット '{dataset_name}' は save_to_disk に対応していません。"
 
-    # データを保存
-    dataset.save_to_disk(temp_dir)
+    temp_dir = f"temp_dataset_{uuid.uuid4().hex}"
+    os.makedirs(temp_dir, exist_ok=True)
 
-    # フォルダパスを指定
-    folder_path = os.path.join(temp_dir, folder_name)
-    if not os.path.exists(folder_path):
-        shutil.rmtree(temp_dir)
-        return f"指定されたフォルダ {folder_name} が見つかりませんでした。"
+    try:
+        dataset.save_to_disk(temp_dir)
+        folder_path = os.path.join(temp_dir, folder_name)
+
+        if not os.path.exists(folder_path):
+            available_folders = os.listdir(temp_dir)
+            return f"指定されたフォルダ '{folder_name}' は見つかりません。利用可能なフォルダ: {available_folders}"
 
-    # ZIPファイルに圧縮
-    with zipfile.ZipFile(output_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, dirs, files in os.walk(folder_path):
-            for file in files:
-                file_path = os.path.join(root, file)
-                arcname = os.path.relpath(file_path, start=folder_path)
-                zipf.write(file_path, arcname)
+        with zipfile.ZipFile(output_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, dirs, files in os.walk(folder_path):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, start=folder_path)
+                    zipf.write(file_path, arcname)
 
-    # 一時ディレクトリを削除
-    shutil.rmtree(temp_dir)
+    except Exception as e:
+        return f"エラーが発生しました: {str(e)}"
+    finally:
+        shutil.rmtree(temp_dir)
 
     return output_zip
 
 def gradio_interface(dataset_name, folder_name):
-    """
-    Gradioインターフェース用関数。
-
-    Args:
-        dataset_name (str): データセット名。
-        folder_name (str): フォルダ名。
-
-    Returns:
-        file: ZIPファイル。
-    """
     output_zip = "output.zip"
     result = zip_folder(dataset_name, folder_name, output_zip)
 
-    if not os.path.exists(result):
-        return f"エラー: {result}"
+    if not os.path.exists(output_zip):
+        return result
 
-    return result
+    return output_zip
 
-# Gradio UI
 description = "Hugging Faceデータセット内の指定フォルダをZIP形式でダウンロード"
 iface = gr.Interface(
     fn=gradio_interface,
@@ -80,7 +56,10 @@ iface = gr.Interface(
         gr.Textbox(label="Hugging Faceデータセット名", placeholder="例: squad"),
         gr.Textbox(label="フォルダ名", placeholder="例: train"),
     ],
-    outputs=gr.File(label="ダウンロード: ZIPファイル"),
+    outputs=[
+        gr.Textbox(label="結果"),  # エラーメッセージ用
+        gr.File(label="ダウンロード: ZIPファイル"),
+    ],
     title="Hugging FaceフォルダZIP変換",
     description=description
 )