import gradio as gr import pandas as pd from datasets import Dataset from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments import torch import os import matplotlib.pyplot as plt def train_model(data_file, model_name, epochs, batch_size, learning_rate): try: if not data_file.name.endswith('.csv'): return "Invalid file format. Please upload a CSV file.", None df = pd.read_csv(data_file.name) if 'prompt' not in df.columns or 'description' not in df.columns: return "CSV file must contain 'prompt' and 'description' columns.", None df['text'] = df['prompt'] + ': ' + df['description'] dataset = Dataset.from_pandas(df[['text']]) 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]'}) tokenizer.pad_token = tokenizer.eos_token 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, remove_columns=['text']) training_args = TrainingArguments( output_dir="./results", 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=10_000, save_total_limit=2, evaluation_strategy="steps", eval_steps=500, load_best_model_at_end=True, metric_for_best_model="eval_loss" ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_datasets, eval_dataset=tokenized_datasets, ) trainer.train() eval_results = trainer.evaluate() model.save_pretrained('./fine-tuned-gpt2') tokenizer.save_pretrained('./fine-tuned-gpt2') train_loss = [log['loss'] for log in trainer.state.log_history if 'loss' in log] eval_loss = [log['eval_loss'] for log in trainer.state.log_history if 'eval_loss' in log] 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('./training_eval_loss.png') return "Training completed successfully.", './training_eval_loss.png' except Exception as e: return f"An error occurred: {str(e)}", None def generate_text(prompt, temperature, top_k, max_length, repetition_penalty, use_comma): 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") input_ids = inputs.input_ids attention_mask = inputs.attention_mask outputs = model.generate( input_ids, attention_mask=attention_mask, max_length=int(max_length), temperature=float(temperature) if temperature > 0 else None, top_k=int(top_k) if top_k > 0 else None, repetition_penalty=float(repetition_penalty), num_return_sequences=1, pad_token_id=tokenizer.eos_token_id, do_sample=True # Assurer que do_sample est activé pour temperature et top_k ) return tokenizer.decode(outputs[0], skip_special_tokens=True) except Exception as e: return f"An error occurred: {str(e)}" def set_preset(preset): if preset == "Default": return 5, 8, 3e-5 elif preset == "Fast Training": return 3, 16, 5e-5 elif preset == "High Accuracy": return 10, 4, 1e-5 with gr.Blocks() as ui: gr.Markdown("# Fine-Tuner | Photographer Alpha7") with gr.Tab("Train Model"): with gr.Row(): data_file = gr.File(label="Upload Data File (CSV)") model_name = gr.Textbox(label="Model Name", value="gpt2") with gr.Row(): preset = gr.Radio(["Default", "Fast Training", "High Accuracy"], label="Preset") epochs = gr.Number(label="Epochs", value=5) batch_size = gr.Number(label="Batch Size", value=8) learning_rate = gr.Number(label="Learning Rate", value=3e-5) preset.change(set_preset, preset, [epochs, batch_size, learning_rate]) train_button = gr.Button("Train Model") train_output = gr.Textbox(label="Training Output") train_graph = gr.Image(label="Training and Validation Loss Graph") train_button.click(train_model, inputs=[data_file, model_name, epochs, batch_size, learning_rate], outputs=[train_output, train_graph]) with gr.Tab("Generate Text"): with gr.Column(): with gr.Row(): generated_text = gr.Textbox(label="Result view") with gr.Row(): prompt = gr.Textbox(label="Prompt") generate_button = gr.Button("Generate Text") with gr.Column(): temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7) top_k = gr.Slider(label="Top K", minimum=1, maximum=100, value=50) max_length = gr.Slider(label="Max Length", minimum=10, maximum=1024, value=128) repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.2) use_comma = gr.Checkbox(label="Use Comma", value=True) generate_button.click(generate_text, inputs=[prompt, temperature, top_k, max_length, repetition_penalty, use_comma], outputs=generated_text) ui.launch()