import os
import zipfile
from huggingface_hub import hf_hub_download
import gradio as gr

# Hugging Faceからデータセットをダウンロードする関数
def download_and_zip(repo_id, folder_name):
    # Hugging Face Hubから対象ファイルをダウンロード
    dataset_path = hf_hub_download(repo_id=repo_id)
    
    folder_path = os.path.join(dataset_path, folder_name)

    # ZIP圧縮の保存先パス
    zip_path = f"{folder_name}.zip"

    # フォルダが存在するかをチェック
    if os.path.exists(folder_path) and os.path.isdir(folder_path):
        # フォルダをZIP圧縮する
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, dirs, files in os.walk(folder_path):
                for file in files:
                    zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
        return zip_path
    else:
        return "指定されたフォルダが見つかりません"

# Gradioインターフェースの作成
def interface(repo_id, folder_name):
    return download_and_zip(repo_id, folder_name)

# Gradioインターフェースのセットアップ
gr.Interface(
    fn=interface,
    inputs=["text", "text"],
    outputs="file",
    live=True,
    description="Hugging Faceのデータセットの指定フォルダをZIPに圧縮します。"
).launch()