# Installing More Dependencies
import torch
from datasets import load_dataset, Dataset
from peft import LoraConfig, AutoPeftModelForCausalLM, get_peft_model, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from trl import SFTTrainer
import os

# Configuration
#MODEL_ID = "NousResearch/Meta-Llama-3.1-8B-Instruct"
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
OUTPUT_DIR = "./fine_tunned_dodel_ul2"
DATASET_NAME = "your_dataset"  # Remplacer par votre dataset
# Charger le modèle et le tokenizer

bnb_config = BitsAndBytesConfig(
      load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype="float16", bnb_4bit_use_double_quant=True
)
try:
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
        quantization_config=bnb_config,
        #max_memory={0: "18GB", "cpu": "24GB"}  # Ajuste la mémoire GPU et CPU
    )
except torch.cuda.OutOfMemoryError:
    print("Erreur de mémoire GPU. Tentative avec des paramètres réduits...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
        quantization_config=bnb_config,
        #max_memory={0: "8GB", "cpu": "32GB"},  # Réduit la mémoire GPU, augmente CPU
        load_in_8bit=True  # Active la quantification 8-bit
    )
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token  # Configure le pad_token comme étant l'eos_token
lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
# Appliquer LoRA au modèle
model = get_peft_model(model, lora_config)
# Préparer le dataset
dataset = dataset = load_dataset(
    "json", 
        data_files={
            "train": "datasets/train.json", 
            "validation": "datasets/validation.json", 
            "test": "datasets/test.json"
        }
)

def tokenize_function(examples):
    # return self.tokenizer(
    #     examples["text"],
    #     padding="max_length",
    #     truncation=True,
    #     max_length=self.config.max_length
    # )
    """
    Fonction de prétraitement des exemples :
    - Concatène le contexte et la question pour créer les entrées
    - Tokenise les entrées et les réponses cibles
    """
    # Concaténer contexte et question pour chaque exemple
    inputs = [context + " " + question for context, question in zip(examples["context"], examples["question"])]
    
    # Extraire les réponses ciblées
    targets = examples["response"]
    
    # Tokenisation des entrées avec padding et troncature
    model_inputs = tokenizer(
        inputs,
        padding="max_length",
        truncation=True,
        max_length=512,
        return_tensors="pt"
    )
    
    # Tokenisation des cibles avec padding et troncature
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(
            targets,
            padding="max_length", 
            truncation=True,
            max_length=128,
            return_tensors="pt"
        )
    
    # Activation des gradients uniquement pour input_ids et labels
    model_inputs["input_ids"] = model_inputs["input_ids"]
    # model_inputs["labels"] = model_inputs["labels"]

    # model_inputs["input_ids"] = model_inputs["input_ids"].cpu().requires_grad_(True).to(model.device)
    # model_inputs["labels"] = model_inputs["labels"].cpu().requires_grad_(True).to(model.device)
    return model_inputs        
        
train_data = dataset['train']
eval_data = dataset['validation']            
train_dataset = train_data.map(tokenize_function, batched=True)
eval_dataset = eval_data.map(tokenize_function, batched=True) if eval_data else None
# tokenized_dataset = dataset.map(
#     lambda x: tokenizer(x["text"], truncation=True, padding="max_length", max_length=512),
#     batched=True
# )
# Configuration de l'entraînement
peft_config = LoraConfig(
    r=8, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
)
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,
    per_device_train_batch_size=4,  # Réduit de 4 à 2
    gradient_accumulation_steps=16,   # Augmenté de 4 à 8
    learning_rate=2e-4,
    fp16=True,
    save_strategy="epoch",
    gradient_checkpointing=True,     # Active le gradient checkpointing
    max_grad_norm=0.3,
    push_to_hub=True                # Limite le gradient pour économiser la mémoire
)
# Initialiser le trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    #dataset_text_field="text",
    args=training_args,
    tokenizer=tokenizer,
    #packing=False,
    max_seq_length=1024
)
# Entraînement
#trainer.train()
try:
    # S'assurer que le modèle est en mode entraînement
    model.train()
    # Activer les gradients
    torch.set_grad_enabled(True)
    trainer.train()

    # Merge avec le modèle original et sauvegarder
    if isinstance(model, PeftModel):
        merged_model = model.merge_and_unload()
    else:
        merged_model = model
        # raise ValueError("Model is not a PeftModel with LoRA adapters")
    # merged_model = model.merge_and_unload()
    merged_model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    # Push sur Hugging Face Hub
    merged_model.push_to_hub(f"{os.getenv('HF_USERNAME')}/Meta-Llama-3.1-8B-Instruct-finetuned")
    tokenizer.push_to_hub(f"{os.getenv('HF_USERNAME')}/Meta-Llama-3.1-8B-Instruct-finetuned")
# except torch.cuda.OutOfMemoryError:
    # print("Erreur de mémoire GPU pendant l'entraînement. Tentative avec des paramètres réduits...")
    # training_args.per_device_train_batch_size //= 2
    # training_args.gradient_accumulation_steps *= 2
    # print(f"Nouveaux paramètres : batch_size={training_args.per_device_train_batch_size}, grad_accum={training_args.gradient_accumulation_steps}")
    # trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset)
    # trainer.train()

except RuntimeError as e:
    print(f"Erreur d'entraînement : {str(e)}")
    raise