File size: 5,566 Bytes
bc63b44 194731c c0d76c2 a1b0975 194731c 038610e a6cb7c8 194731c c0d76c2 bc63b44 038610e 4d120d4 f9e951b a6cb7c8 194731c 0958d38 194731c 0958d38 d06ead9 194731c 1c20c42 194731c 1c20c42 8bef298 a1b0975 e585d7a a1b0975 8bef298 77c210a 8bef298 194731c 1c20c42 a1b0975 194731c 22093a9 0314370 c0d76c2 3c24b96 c0d76c2 194731c a1b0975 194731c 6c1ed42 0958d38 4d120d4 3c24b96 194731c 5c817b9 194731c a1b0975 9a0b862 194731c c17c736 194731c c17c736 038610e 3c24b96 0314370 57918ff 3c24b96 194731c 172d00c 194731c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import spaces
import gradio as gr
from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM, TrainerCallback
from transformers import DataCollatorForSeq2Seq
from datasets import load_dataset
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):
# Assuming 'text' is the input and 'target' is the expected output
model_inputs = tokenizer(
examples['text'],
max_length=max_length, # Set to None for dynamic padding
padding=False, # Disable padding here, we will handle it later
truncation=True,
)
# Setup the decoder input IDs (shifted right)
labels = tokenizer(
examples['target'],
max_length=max_length, # Set to None for dynamic padding
padding=False, # Disable padding here, we will handle it later
truncation=True,
text_target=examples['target'] # Use text_target for target text
)
# Add labels to the model inputs
model_inputs["labels"] = labels["input_ids"]
return model_inputs
tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=16)
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
# 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=data_collator['train'],
eval_dataset=data_collator['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()}")
|