gradio-3 / app.py
Kevin Fink
init
75c24c0
raw
history blame
5.96 kB
import spaces
import gradio as gr
from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM, TrainerCallback
from transformers import DataCollatorForSeq2Seq
from datasets import load_dataset
from datasets import concatenate_datasets
import traceback
from huggingface_hub import login
from peft import get_peft_model, LoraConfig
class LoggingCallback(TrainerCallback):
def on_step_end(self, args, state, control, kwargs):
# Log the learning rate
current_lr = state.optimizer.param_groups[0]['lr']
print(f"Current Learning Rate: {current_lr}")
def on_epoch_end(self, args, state, control, kwargs):
# Log the error rate (assuming you have a metric to calculate it)
# Here we assume you have a way to get the validation loss
if state.best_metric is not None:
error_rate = 1 - state.best_metric # Assuming best_metric is accuracy
print(f"Current Error Rate: {error_rate:.4f}")
@spaces.GPU
def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
try:
login(api_key.strip())
lora_config = LoraConfig(
r=16, # Rank of the low-rank adaptation
lora_alpha=32, # Scaling factor
lora_dropout=0.1, # Dropout for LoRA layers
bias="none" # Bias handling
)
# Load the dataset
dataset = load_dataset(dataset_name.strip())
# Load the model and tokenizer
model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
#model = get_peft_model(model, lora_config)
tokenizer = AutoTokenizer.from_pretrained(model_name)
chunk_size = 1000
max_length = 128
# Tokenize the dataset
def tokenize_function(examples):
# Tokenize the input text
model_inputs = tokenizer(
examples['text'],
max_length=max_length,
padding=True,
truncation=True,
)
# Tokenize the target text
labels = tokenizer(
examples['target'],
max_length=max_length,
padding=True,
truncation=True,
)
# Add labels to the model inputs
model_inputs["labels"] = labels["input_ids"]
return model_inputs
# Define a function to process the dataset in chunks
def tokenize_in_chunks(dataset, chunk_size=1000):
tokenized_datasets = []
for i in range(0, len(dataset), chunk_size):
chunk = dataset[i:i + chunk_size]
tokenized_chunk = chunk.map(tokenize_function, batched=True)
tokenized_datasets.append(tokenized_chunk)
# Concatenate all tokenized chunks into a single dataset
return tokenized_datasets
# Tokenize the dataset in chunks
tokenized_datasets = tokenize_in_chunks(dataset, chunk_size=1000)
# If you want to combine all chunks into a single dataset
final_tokenized_dataset = concatenate_datasets(tokenized_datasets)
# Set training arguments
training_args = TrainingArguments(
output_dir='./results',
eval_strategy="epoch",
save_strategy='epoch',
learning_rate=lr*0.000001,
per_device_train_batch_size=int(batch_size),
per_device_eval_batch_size=1,
num_train_epochs=int(num_epochs),
weight_decay=0.01,
gradient_accumulation_steps=int(grad),
load_best_model_at_end=True,
metric_for_best_model="accuracy",
greater_is_better=True,
logging_dir='./logs',
logging_steps=10,
#push_to_hub=True,
hub_model_id=hub_id.strip(),
fp16=True,
#lr_scheduler_type='cosine',
)
# Create Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=final_tokenized_dataset['train'],
eval_dataset=final_tokenized_dataset['test'],
#callbacks=[LoggingCallback()],
)
# Fine-tune the model
trainer.train()
trainer.push_to_hub(commit_message="Training complete!")
except Exception as e:
return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
return 'DONE!'#model
'''
# Define Gradio interface
def predict(text):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
outputs = model(inputs)
predictions = outputs.logits.argmax(dim=-1)
return "Positive" if predictions.item() == 1 else "Negative"
'''
# Create Gradio interface
try:
iface = gr.Interface(
fn=fine_tune_model,
inputs=[
gr.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
gr.Textbox(label="HF hub to push to after training"),
gr.Textbox(label="HF API token"),
gr.Slider(minimum=1, maximum=10, value=3, label="Number of Epochs", step=1),
gr.Slider(minimum=1, maximum=16, value=1, label="Batch Size", step=1),
gr.Slider(minimum=1, maximum=1000, value=1, label="Learning Rate (e-6)", step=1),
gr.Slider(minimum=1, maximum=100, value=1, label="Gradient accumulation (e-1)", step=1),
],
outputs="text",
title="Fine-Tune Hugging Face Model",
description="This interface allows you to fine-tune a Hugging Face model on a specified dataset."
)
# Launch the interface
iface.launch()
except Exception as e:
print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")