File size: 1,446 Bytes
51027b7 | 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 | #!/usr/bin/env python3
"""
Bengali Sentiment - Fine-tuning Script
=======================================
Usage: python fine_tune_script.py
"""
from datasets import load_dataset
from transformers import (
AutoTokenizer, AutoModelForSequenceClassification,
Trainer, TrainingArguments, DataCollatorWithPadding
)
DATASET = "Swarnadeep-28/bn_code_mix_sentiment_dataset"
MODEL = "distilbert-base-multilingual-cased"
LABEL_MAP = {"Negative": 0, "Neutral": 1, "Positive": 2, "Mixed": 3}
def preprocess(examples):
texts = [str(t) for t in examples["text"]]
labels = [LABEL_MAP[l] for l in examples["label"]]
tok = AutoTokenizer.from_pretrained(MODEL)
tokenized = tok(texts, truncation=True, padding=True, max_length=128)
tokenized["labels"] = labels
return tokenized
ds = load_dataset(DATASET)
train = ds["train"].select(range(2000)).map(preprocess, batched=True, remove_columns=ds["train"].column_names)
val = ds["validation"].select(range(500)).map(preprocess, batched=True, remove_columns=ds["validation"].column_names)
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=4)
args = TrainingArguments(output_dir="./output", num_train_epochs=3, per_device_train_batch_size=8, eval_strategy="epoch")
trainer = Trainer(model=model, args=args, train_dataset=train, eval_dataset=val, data_collator=DataCollatorWithPadding())
trainer.train()
trainer.save_model("./bengali_sentiment_finetuned")
|