app.py

import gradio as gr
import pandas as pd
from datasets import load_dataset, Dataset
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments
import torch
import os
import matplotlib.pyplot as plt
from huggingface_hub import HfApi  # ここを修正しました
import json
import io
from datetime import datetime

# グローバル変数で検出された列を保存
columns = []

# ファイル読み込み関数
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='訓練損失')
        plt.plot(eval_loss, label='評価損失')
        plt.xlabel('ステップ数')
        plt.ylabel('損失')
        plt.title('訓練と評価の損失')
        plt.legend()
        plt.savefig(os.path.join(output_dir, 'training_eval_loss.png'))

        # モデルのHuggingFaceにアップロード
        hf_api = HfApi()
        hf_api.upload_folder(
            folder_path=output_dir,
            path_in_repo=".",
            repo_id=model_name,
            token=hf_token
        )

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

# テキスト生成関数
def generate_text(prompt, temperature, top_k, top_p, max_length, repetition_penalty, use_comma, batch_size):
    try:
        model_name = "./fine-tuned-gpt2"
        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        model = GPT2LMHeadModel.from_pretrained(model_name)

        if use_comma:
            prompt = prompt.replace('.', ',')

        inputs = tokenizer(prompt, return_tensors="pt", padding=True)
        attention_mask = inputs.attention_mask
        outputs = model.generate(
            inputs.input_ids,
            attention_mask=attention_mask,
            max_length=int(max_length),
            temperature=float(temperature),
            top_k=int(top_k),
            top_p=float(top_p),
            repetition_penalty=float(repetition_penalty),
            num_return_sequences=int(batch_size),
            pad_token_id=tokenizer.eos_token_id
        )

        return [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
    except Exception as e:
        return f"エラーが発生しました: {str(e)}"

# 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()