File size: 5,960 Bytes
bc63b44 194731c c0d76c2 a1b0975 194731c 5a28d7d 194731c 038610e a6cb7c8 194731c c0d76c2 bc63b44 038610e 4d120d4 f9e951b a6cb7c8 194731c 0958d38 194731c 0958d38 d06ead9 194731c 1c20c42 194731c 5a28d7d a1b0975 5a28d7d a1b0975 8bef298 5a28d7d 77c210a 5a28d7d 77c210a 8bef298 5a28d7d 75c24c0 5a28d7d 194731c 22093a9 0314370 c0d76c2 3c24b96 c0d76c2 194731c a1b0975 194731c 6c1ed42 0958d38 4d120d4 3c24b96 194731c 5c817b9 194731c 5a28d7d 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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
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()}")
|