File size: 3,397 Bytes
5089144 dc44a31 5089144 dc44a31 5089144 dc44a31 5089144 7308396 5089144 |
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 |
#!/usr/bin/env python3
import torch
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
DataCollatorWithPadding
)
from datasets import Dataset
import json
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class AcoliTrainer:
def __init__(self, model_name="xlm-roberta-base", num_labels=3):
self.model_name = model_name
self.num_labels = num_labels
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(
model_name,
num_labels=num_labels
)
def load_data(self, jsonl_path):
"""Load data from JSONL file"""
texts = []
labels = []
with open(jsonl_path, 'r', encoding='utf-8') as f:
for line in f:
data = json.loads(line)
texts.append(data['text'])
labels.append(data['label'])
return Dataset.from_dict({
'text': texts,
'label': labels
})
def preprocess_function(self, examples):
"""Tokenize the texts"""
return self.tokenizer(
examples['text'],
truncation=True,
padding=True,
max_length=512
)
def train(self, train_path, output_dir="./acoli-model"):
"""Train the model"""
# Load and preprocess data
logger.info("Loading training data...")
dataset = self.load_data(train_path)
tokenized_dataset = dataset.map(self.preprocess_function, batched=True)
# Split dataset (80% train, 20% validation)
train_test_split = tokenized_dataset.train_test_split(test_size=0.2)
train_dataset = train_test_split['train']
eval_dataset = train_test_split['test']
# Training arguments
training_args = TrainingArguments(
output_dir=output_dir,
learning_rate=2e-5,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=3,
weight_decay=0.01,
evaluation_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
push_to_hub=False, # Set to True if you want to push to HF Hub
)
# Data collator
data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
# Initialize Trainer
trainer = Trainer(
model=self.model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=self.tokenizer,
data_collator=data_collator,
)
# Start training
logger.info("Starting training...")
trainer.train()
# Save the model
logger.info(f"Saving model to {output_dir}")
trainer.save_model(output_dir)
self.tokenizer.save_pretrained(output_dir)
return trainer
if __name__ == "__main__":
# Example usage
trainer = AcoliTrainer()
# Train the model
trained_trainer = trainer.train("path/to/your/data.jsonl")
print("Training completed successfully!") |