%%capture
!pip install transformers accelerate bitsandbytes

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline, AutoConfig
import torch

model_name = "vakodiya/Llama-3-8B-instruct-4bit-salesbot"
config = AutoConfig.from_pretrained(model_name)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
# Model evaluation mode
model.eval()

def Trained_Llama3_1_inference(prompt):
    model.eval()
    conversation=[
        {"role": "user", "content": prompt},
    ]
    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt", padding=True, truncation=True, return_attention_mask=True)
    if input_ids.shape[1] > 8192:
        input_ids = input_ids[:, -8192:]
        return "Input tokens more than 8k"
    inputs = input_ids.to(model.device)
    attention_mask = torch.ones_like(inputs, dtype=torch.long)
    final_prompt=tokenizer.decode(inputs[0])
    outputs = model.generate(inputs, max_new_tokens=256, temperature=0.4,attention_mask=attention_mask,pad_token_id=tokenizer.pad_token_id)
    response = tokenizer.decode(outputs[0])
    final_response= response.replace(final_prompt,"").replace('<|eot_id|>',"")  # Exclude prompt from response
    return final_response

%%capture
!pip install -U transformers[torch] datasets
!pip install -q bitsandbytes trl peft accelerate
!pip install flash-attn --no-build-isolation
!pip install huggingface_hub

from transformers import BitsAndBytesConfig,  AutoTokenizer,  AutoModelForCausalLM, TrainingArguments
from trl import SFTTrainer
from peft import LoraConfig
from huggingface_hub import notebook_login
import torch
from huggingface_hub import login
from datasets import Dataset
from kaggle_secrets import UserSecretsClient
import os

hf_token = UserSecretsClient().get_secret("HF_TOKEN_LLAMA3")
login(token = hf_token)
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Use only GPU 0

Data_examples = [{"instruction":"Who has taken oath as Prime minister of India in 2024", "response":" Shri Narendra Modi has taken oath as Prime minister of india on 9th June 2024. He is now become prime minister having 3 consecutive terms."},
                  ...................................................................................,]

processed_data = []
for example in Data_examples :    
    processed_data.append({'text':f"{example['instruction']} \n {example['response']}"})

# Create a Dataset from the list of dictionaries
dataset = Dataset.from_list(processed_data)

# Split into train and test Data sets

dataset = dataset.train_test_split(test_size=0.01)
# Access train and test splits

train_dataset = dataset['train']
test_dataset = dataset['test']

model_path="/kaggle/input/llama-3.1/transformers/8b-instruct/2"  # Change it according to your model path in Notebook
trained_model_name = "Llama-3-8B-instruct-4bit-finetuned"
output_dir = 'kaggle/working/' + trained_model_name 

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,)

peft_config = LoraConfig(
        r=16,
        lora_alpha=16,
        lora_dropout=0.1,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
)

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    quantization_config=quantization_config,
    device_map="auto")

# Use eos_token as pad_token
tokenizer.pad_token = tokenizer.eos_token

training_args = TrainingArguments(
    fp16=False, # specify bf16=True instead when training on GPUs that support bf16 else fp16
    bf16=True,
    do_eval=True,
    eval_strategy="epoch",
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    learning_rate=2.0e-05,
    log_level="info",
    logging_steps=5,
    logging_strategy="steps",
    lr_scheduler_type="cosine",
    max_steps=-1,
    num_train_epochs=1,       # Number of times training will go through with same dataset.
    output_dir=output_dir,
    overwrite_output_dir=True,
    per_device_eval_batch_size=8, #  You can reduce if out-of memory errors occurs
    per_device_train_batch_size=8, #  You can reduce if out-of memory errors occurs
    report_to="none",  # for skipping wandb logging
    save_strategy="no",
    save_total_limit=None,
)

trainer = SFTTrainer(
        model=model,                   # Use above quantized model
        args=training_args,
        train_dataset=train_dataset,          # If Training Fails Try to reduce Dataset Size
        eval_dataset=test_dataset,
        dataset_text_field="text",
        tokenizer=tokenizer,
        packing=False,             # Setting it True will Reduce dataset size as it will exclude similar examples occuring repetitive
        peft_config=peft_config,
        max_seq_length=1024,
    )

# To clear out cache for unsuccessful run
torch.cuda.empty_cache()

train_result = trainer.train()

merged_model = trainer.model.merge_and_unload()
merged_model.save_pretrained("merged_model",safe_serialization=True)
tokenizer.save_pretrained("merged_model")

merged_model.push_to_hub("username/model_name")
tokenizer.push_to_hub("username/model_name")