import os
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from datasets import load_dataset

def load_model_and_tokenizer(model_name):
    """
    Load the model and tokenizer.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return model, tokenizer

def load_and_tokenize_dataset(dataset_name, tokenizer, max_length=512):
    """
    Load and tokenize the dataset.
    """
    dataset = load_dataset(dataset_name)

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

    tokenized_datasets = dataset.map(tokenize_function, batched=True)
    return tokenized_datasets

def setup_training_args(output_dir="./results", per_device_train_batch_size=2, per_device_eval_batch_size=2,
                        gradient_accumulation_steps=8, num_train_epochs=3, learning_rate=5e-5, weight_decay=0.01,
                        warmup_steps=500, logging_steps=100, fp16=True):
    """
    Set up training arguments.
    """
    training_args = TrainingArguments(
        output_dir=output_dir,
        evaluation_strategy="epoch",
        per_device_train_batch_size=per_device_train_batch_size,
        per_device_eval_batch_size=per_device_eval_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        num_train_epochs=num_train_epochs,
        save_strategy="epoch",
        save_total_limit=2,
        logging_dir="./logs",
        logging_steps=logging_steps,
        report_to="none",
        fp16=fp16,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        warmup_steps=warmup_steps,
        dataloader_num_workers=4,
        push_to_hub=False
    )
    return training_args

def save_model_and_tokenizer(model, tokenizer, save_dir):
    """
    Save the model and tokenizer.
    """
    os.makedirs(save_dir, exist_ok=True)
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)
    print(f"Model and tokenizer saved at {save_dir}")