app.py

import gradio as gr
import pandas as pd
from datasets import Dataset, load_dataset
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, HfApi
import torch
import os
import matplotlib.pyplot as plt
import json
import io
from datetime import datetime

# Variables globales pour stocker les colonnes détectées
columns = []

# Hugging Faceにアクセスするためのアクセストークン
hf_token = "YOUR_HUGGINGFACE_ACCESS_TOKEN"

# ファイル読み込み機能
def read_file(data_file):
    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ファイルをアップロードしてください。"

        # 列を検出
        columns = df.columns.tolist()
        return columns
    except Exception as e:
        return f"エラーが発生しました: {str(e)}"


# 列のバリデーション
def validate_columns(prompt_col, description_col):
    if prompt_col not in columns or description_col not in columns:
        return False
    return True


# モデルの訓練
def train_model(data_file, 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)

        # データのプレビュー
        preview = df.head().to_string(index=False)

        # トレーニングテキストの準備
        df['text'] = df[prompt_col] + ': ' + df[description_col]
        dataset = Dataset.from_pandas(df[['text']])

        # GPT-2トークナイザーとモデルの初期化
        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        model = GPT2LMHeadModel.from_pretrained(model_name)

        # パディングトークンの追加
        if tokenizer.pad_token is None:
            tokenizer.add_special_tokens({'pad_token': '[PAD]'})
            model.resize_token_embeddings(len(tokenizer))

        # データのトークナイズ
        def tokenize_function(examples):
            tokens = tokenizer(examples['text'], padding="max_length", truncation=True, max_length=128)
            tokens['labels'] = tokens['input_ids'].copy()
            return tokens

        tokenized_datasets = dataset.map(tokenize_function, batched=True)

        # ハイパーパラメータの設定
        training_args = TrainingArguments(
            output_dir=output_dir,
            overwrite_output_dir=True,
            num_train_epochs=int(epochs),
            per_device_train_batch_size=int(batch_size),
            per_device_eval_batch_size=int(batch_size),
            warmup_steps=1000,
            weight_decay=0.01,
            learning_rate=float(learning_rate),
            logging_dir="./logs",
            logging_steps=10,
            save_steps=500,
            save_total_limit=2,
            evaluation_strategy="steps",
            eval_steps=500,
            load_best_model_at_end=True,
            metric_for_best_model="eval_loss"
        )

        # Trainerの設定
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_datasets,
            eval_dataset=tokenized_datasets,
        )

        # 訓練と評価
        trainer.train()
        eval_results = trainer.evaluate()

        # Fine-tunedモデルの保存
        model.save_pretrained(output_dir)
        tokenizer.save_pretrained(output_dir)

        # トレーニングと評価の損失グラフ生成
        train_loss = [x['loss'] for x in trainer.state.log_history if 'loss' in x]
        eval_loss = [x['eval_loss'] for x in trainer.state.log_history if 'eval_loss' in x]
        plt.plot(train_loss, label='Training Loss')
        plt.plot(eval_loss, label='Validation Loss')
        plt.xlabel('Steps')
        plt.ylabel('Loss')
        plt.title('Training and Validation Loss')
        plt.legend()
        plt.savefig(os.path.join(output_dir, 'training_eval_loss.png'))

        # Hugging Faceにアップロード
        upload_response = upload_model_to_huggingface(output_dir, model_name, hf_token)

        return f"訓練が成功しました。\nデータプレビュー:\n{preview}", eval_results, upload_response
    except Exception as e:
        return f"エラーが発生しました: {str(e)}"


# モデルをHugging Faceにアップロード
def upload_model_to_huggingface(output_dir, model_name, hf_token):
    try:
        api = HfApi()
        repo_url = api.create_repo(model_name, exist_ok=True)  # リポジトリが既にあればそのまま使用
        api.upload_folder(
            folder_path=output_dir,
            repo_id=model_name,
            path_in_repo=".",
            use_auth_token=hf_token
        )
        return f"モデルがHugging Faceに正常にアップロードされました。\nリポジトリURL: https://huggingface.co/{model_name}"
    except Exception as e:
        return f"モデルのアップロード中にエラーが発生しました: {str(e)}"


# UI設定
def generate_text(prompt, temperature, top_k, top_p, max_length, repetition_penalty, use_comma, batch_size):
    # 生成ロジック（実際のモデル使用コードを挿入）
    return "生成されたテキスト"

# UI設定
with gr.Blocks() as ui:
    with gr.Row():
        data_file = gr.File(label="データファイル", file_types=[".csv", ".json", ".xlsx"])
        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")
        hf_token = gr.Textbox(label="Hugging Face アクセストークン")

    with gr.Row():
        validate_button = gr.Button("列検証")
        output = gr.Textbox(label="出力")

    validate_button.click(
        read_file, 
        inputs=[data_file],
        outputs=[output]
    )

    with gr.Row():
        train_button = gr.Button("訓練開始")
        result_output = gr.Textbox(label="訓練結果", lines=20)

    train_button.click(
        train_model,
        inputs=[data_file, model_name, epochs, batch_size, learning_rate, output_dir, prompt_col, description_col, hf_token],
        outputs=[result_output]
    )

ui.launch()