|
--- |
|
datasets: |
|
- asr-malayalam/indicvoices-v1a |
|
- Tensoic/GPTeacher-Malayalam |
|
language: |
|
- ml |
|
- en |
|
metrics: |
|
- accuracy |
|
base_model: |
|
- deepseek-ai/DeepSeek-R1 |
|
pipeline_tag: translation |
|
--- |
|
import os |
|
import argparse |
|
import pandas as pd |
|
from datasets import Dataset |
|
from transformers import ( |
|
AutoTokenizer, |
|
AutoModelForSeq2SeqLM, |
|
Seq2SeqTrainingArguments, |
|
Seq2SeqTrainer, |
|
DataCollatorForSeq2Seq |
|
) |
|
from utils import compute_metrics |
|
|
|
def load_dataset(file_path): |
|
"""Load and prepare the dataset.""" |
|
df = pd.read_csv(file_path) |
|
dataset = Dataset.from_pandas(df) |
|
# Split dataset into train and validation |
|
split_dataset = dataset.train_test_split(test_size=0.1) |
|
return split_dataset |
|
|
|
def preprocess_function(examples, tokenizer, max_length=128): |
|
"""Tokenize the texts.""" |
|
inputs = [ex for ex in examples["english_text"]] |
|
targets = [ex for ex in examples["malayalam_text"]] |
|
|
|
model_inputs = tokenizer( |
|
inputs, |
|
max_length=max_length, |
|
truncation=True, |
|
padding="max_length", |
|
) |
|
|
|
with tokenizer.as_target_tokenizer(): |
|
labels = tokenizer( |
|
targets, |
|
max_length=max_length, |
|
truncation=True, |
|
padding="max_length", |
|
) |
|
|
|
model_inputs["labels"] = labels["input_ids"] |
|
return model_inputs |
|
|
|
def main(args): |
|
# Load tokenizer and model |
|
model_name = "google/mt5-small" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) |
|
|
|
# Load and preprocess dataset |
|
dataset = load_dataset("dataset/malayalam_dataset.csv") |
|
|
|
# Tokenize datasets |
|
tokenized_datasets = dataset.map( |
|
lambda x: preprocess_function(x, tokenizer), |
|
batched=True, |
|
remove_columns=dataset["train"].column_names |
|
) |
|
|
|
# Define training arguments |
|
training_args = Seq2SeqTrainingArguments( |
|
output_dir="./model", |
|
evaluation_strategy="epoch", |
|
learning_rate=args.learning_rate, |
|
per_device_train_batch_size=args.batch_size, |
|
per_device_eval_batch_size=args.batch_size, |
|
num_train_epochs=args.epochs, |
|
weight_decay=0.01, |
|
save_total_limit=2, |
|
predict_with_generate=True, |
|
logging_dir="./logs", |
|
logging_steps=100, |
|
push_to_hub=True, |
|
) |
|
|
|
# Create data collator |
|
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model) |
|
|
|
# Initialize trainer |
|
trainer = Seq2SeqTrainer( |
|
model=model, |
|
args=training_args, |
|
train_dataset=tokenized_datasets["train"], |
|
eval_dataset=tokenized_datasets["test"], |
|
data_collator=data_collator, |
|
tokenizer=tokenizer, |
|
compute_metrics=compute_metrics |
|
) |
|
|
|
# Train the model |
|
trainer.train() |
|
|
|
# Save the model |
|
trainer.save_model("./model") |
|
tokenizer.save_pretrained("./model") |
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--epochs", type=int, default=3) |
|
parser.add_argument("--batch_size", type=int, default=8) |
|
parser.add_argument("--learning_rate", type=float, default=2e-5) |
|
args = parser.parse_args() |
|
main(args) |
