Spaces:
Sleeping
Sleeping
File size: 1,517 Bytes
2b3de4f 24cc0f4 2b3de4f d6d26ef 2b3de4f d6d26ef 2b3de4f d6d26ef 2b3de4f 24cc0f4 2b3de4f |
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 |
from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset
# Load dataset (French dataset example: Allociné)
dataset = load_dataset("allocine")
dataset["train"] = dataset["train"].select(range(10)) # Train on 500 samples
dataset["test"] = dataset["test"].select(range(5)) # Test on 200 samples
# Load tokenizer
model_name = "distilbert-base-multilingual-cased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Tokenize data
def tokenize(batch):
return tokenizer(batch["review"], padding="max_length", truncation=True)
dataset = dataset.map(tokenize, batched=True)
# Load model
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
# Training arguments
training_args = TrainingArguments(
output_dir="./models",
evaluation_strategy="epoch", # Ensure this matches the save_strategy
save_strategy="epoch", # Change this to "epoch" to match evaluation_strategy
load_best_model_at_end=True, # Ensures best model is loaded
save_total_limit=2, # Keep only the last 2 models to save space
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=3,
weight_decay=0.01,
)
# Trainer setup
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
)
# Train model
trainer.train()
# Save model
model.save_pretrained("./models")
tokenizer.save_pretrained("./models")
|