Fine-tuning Models

from datasets import load_dataset
from transformers import AutoTokenizer

# Load a dataset from the Hub
dataset = load_dataset("imdb")
print(dataset)
# DatasetDict({
#     train: Dataset({features: ['text', 'label'], num_rows: 25000}),
#     test: Dataset({features: ['text', 'label'], num_rows: 25000})
# })

# Tokenize the dataset
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

from transformers import (
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
)
import numpy as np
from datasets import load_metric

# Load pre-trained model
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

# Define training arguments
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir="./logs",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

# Define evaluation metric
metric = load_metric("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

# Create Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    compute_metrics=compute_metrics,
)

# Train!
trainer.train()

# Evaluate the fine-tuned model
results = trainer.evaluate()
print(f"Accuracy: {results['eval_accuracy']:.4f}")
print(f"Loss: {results['eval_loss']:.4f}")

# Login to Hugging Face
from huggingface_hub import login
login(token="your_token_here")

# Push model and tokenizer to the Hub
trainer.push_to_hub("my-fine-tuned-sentiment-model")

# Or push individually
model.push_to_hub("my-fine-tuned-sentiment-model")
tokenizer.push_to_hub("my-fine-tuned-sentiment-model")