Update app.py

app.py

@@ -13,24 +13,20 @@ from datetime import datetime
 # グローバル変数で検出された列を保存
 columns = []
 
-# ファイル読み込み関数
-def read_file(data_file):
+# データセットをロードする関数
+def load_data(dataset_name):
     global columns
     try:
-        # ファイルをロード
-        file_extension = os.path.splitext(data_file.name)[1]
-        if file_extension == '.csv':
-            df = pd.read_csv(data_file.name)
-        elif file_extension == '.json':
-            df = pd.read_json(data_file.name)
-        elif file_extension == '.xlsx':
-            df = pd.read_excel(data_file.name)
-        else:
-            return "無効なファイル形式です。CSV, JSON, Excelファイルをアップロードしてください。"
-
-        # 列を検出
+        # Hugging Faceのデータセットをロード
+        dataset = load_dataset(dataset_name)
+
+        # 最初のデータをプレビューとして表示
+        df = pd.DataFrame(dataset['train'])
+        
+        # 列名を検出
         columns = df.columns.tolist()
-        return columns
+
+        return columns, df.head().to_string(index=False)
     except Exception as e:
         return f"エラーが発生しました: {str(e)}"
 
@@ -41,27 +37,24 @@ def validate_columns(prompt_col, description_col):
     return True
 
 # モデル訓練関数
-def train_model(data_file, model_name, epochs, batch_size, learning_rate, output_dir, prompt_col, description_col, hf_token):
+def train_model(dataset_name, model_name, epochs, batch_size, learning_rate, output_dir, prompt_col, description_col, hf_token):
     try:
         # 列の検証
         if not validate_columns(prompt_col, description_col):
             return "無効な列選択です。データセット内の列を確認してください。"
 
-        # ファイルのロード
-        file_extension = os.path.splitext(data_file.name)[1]
-        if file_extension == '.csv':
-            df = pd.read_csv(data_file.name)
-        elif file_extension == '.json':
-            df = pd.read_json(data_file.name)
-        elif file_extension == '.xlsx':
-            df = pd.read_excel(data_file.name)
+        # Hugging Faceのデータセットをロード
+        dataset = load_dataset(dataset_name)
+        
+        # 訓練データを取得
+        df = pd.DataFrame(dataset['train'])
 
         # データのプレビュー
         preview = df.head().to_string(index=False)
 
         # 訓練用テキストの準備
         df['text'] = df[prompt_col] + ': ' + df[description_col]
-        dataset = Dataset.from_pandas(df[['text']])
+        train_dataset = Dataset.from_pandas(df[['text']])
 
         # GPT-2のトークナイザーとモデルを初期化
         tokenizer = GPT2Tokenizer.from_pretrained(model_name)
@@ -78,7 +71,7 @@ def train_model(data_file, model_name, epochs, batch_size, learning_rate, output
             tokens['labels'] = tokens['input_ids'].copy()
             return tokens
 
-        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        tokenized_datasets = train_dataset.map(tokenize_function, batched=True)
 
         # 訓練のための設定
         training_args = TrainingArguments(
@@ -171,14 +164,14 @@ def generate_text(prompt, temperature, top_k, top_p, max_length, repetition_pena
 # UI設定
 with gr.Blocks() as ui:
     with gr.Row():
-        data_file = gr.File(label="データファイル", file_types=[".csv", ".json", ".xlsx"])
+        dataset_name = gr.Textbox(label="データセット名", value="imdb")  # ここにデータセット名を入力
         model_name = gr.Textbox(label="モデル名", value="gpt2")
         epochs = gr.Number(label="エポック数", value=3, minimum=1)
         batch_size = gr.Number(label="バッチサイズ", value=4, minimum=1)
         learning_rate = gr.Number(label="学習率", value=5e-5, minimum=1e-7, maximum=1e-2, step=1e-7)
         output_dir = gr.Textbox(label="出力ディレクトリ", value="./output")
-        prompt_col = gr.Textbox(label="プロンプト列名", value="prompt")
-        description_col = gr.Textbox(label="説明列名", value="description")
+        prompt_col = gr.Textbox(label="プロンプト列名", value="text")  # 例：IMDBのレビュー列名
+        description_col = gr.Textbox(label="説明列名", value="label")  # 例：IMDBのラベル列名
         hf_token = gr.Textbox(label="Hugging Face アクセストークン")
 
     with gr.Row():
@@ -186,8 +179,8 @@ with gr.Blocks() as ui:
         output = gr.Textbox(label="出力")
 
     validate_button.click(
-        read_file, 
-        inputs=[data_file],
+        load_data, 
+        inputs=[dataset_name],
         outputs=[output]
     )
 
@@ -197,7 +190,7 @@ with gr.Blocks() as ui:
 
     train_button.click(
         train_model,
-        inputs=[data_file, model_name, epochs, batch_size, learning_rate, output_dir, prompt_col, description_col, hf_token],
+        inputs=[dataset_name, model_name, epochs, batch_size, learning_rate, output_dir, prompt_col, description_col, hf_token],
         outputs=[result_output]
     )