pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

pip install -U huggingface_hub

import os


os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

from unsloth import FastLanguageModel
import torch
max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.
# 4bit pre quantized models we support for 4x faster downloading + no OOMs.
fourbit_models = [
  "unsloth/Meta-Llama-3.1-8B-bnb-4bit",      # Llama-3.1 2x faster
  "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
  "unsloth/Meta-Llama-3.1-70B-bnb-4bit",
  "unsloth/Meta-Llama-3.1-405B-bnb-4bit",    # 4bit for 405b!
  "unsloth/Mistral-Small-Instruct-2409",     # Mistral 22b 2x faster!
  "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
  "unsloth/Phi-3.5-mini-instruct",           # Phi-3.5 2x faster!
  "unsloth/Phi-3-medium-4k-instruct",
  "unsloth/gemma-2-9b-bnb-4bit",
  "unsloth/gemma-2-27b-bnb-4bit",            # Gemma 2x faster!
  "unsloth/Llama-3.2-1B-bnb-4bit",           # NEW! Llama 3.2 models
  "unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
  "unsloth/Llama-3.2-3B-bnb-4bit",
  "unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
] # More models at https://huggingface.co/unsloth


model, tokenizer = FastLanguageModel.from_pretrained(
  model_name = "unsloth/Llama-3.2-3B-Instruct", # or choose "unsloth/Llama-3.2-1B-Instruct"
  max_seq_length = max_seq_length,
  dtype = dtype,
  load_in_4bit = load_in_4bit,
  # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)

model = FastLanguageModel.get_peft_model(
  model,
  r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
  target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
  "gate_proj", "up_proj", "down_proj",],
  lora_alpha = 16,
  lora_dropout = 0, # Supports any, but = 0 is optimized
  bias = "none",    # Supports any, but = "none" is optimized
  # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
  use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
  random_state = 3407,
  use_rslora = False,  # We support rank stabilized LoRA
  loftq_config = None, # And LoftQ
)

from unsloth.chat_templates import get_chat_template


tokenizer = get_chat_template(
    tokenizer,
    chat_template = "llama-3.1",
)


def formatting_prompts_func(examples):
    convos = examples["conversations"]
    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
    return { "text" : texts, }
pass


from datasets import load_dataset
dataset = load_dataset("mlabonne/FineTome-100k", split = "train")

from unsloth.chat_templates import standardize_sharegpt
dataset = standardize_sharegpt(dataset)
dataset = dataset.map(formatting_prompts_func, batched = True,)

dataset[5]["conversations"]

[{'content': 'How do astronomers determine the original wavelength of light emitted by a celestial body at rest, which is necessary for measuring its speed using the Doppler effect?',
'role': 'user'},
{'content': 'Astronomers make use of the unique spectral fingerprints of elements found in stars. These elements emit and absorb light at specific, known wavelengths, forming an absorption spectrum. By analyzing the light received from distant stars and comparing it to the laboratory-measured spectra of these elements, astronomers can identify the shifts in these wavelengths due to the Doppler effect. The observed shift tells them the extent to which the light has been redshifted or blueshifted, thereby allowing them to calculate the speed of the star along the line of sight relative to Earth.',
'role': 'assistant'}]

dataset[5]["text"]

from trl import SFTTrainer
from transformers import TrainingArguments, DataCollatorForSeq2Seq
from unsloth import is_bfloat16_supported


trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
    dataset_num_proc = 2,
    packing = False, # Can make training 5x faster for short sequences.
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        # num_train_epochs = 1, # Set this for 1 full training run.
        max_steps = 60,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
    ),
)

from unsloth.chat_templates import train_on_responses_only
trainer = train_on_responses_only(
    trainer,
    instruction_part = "<|start_header_id|>user<|end_header_id|>\n\n",
    response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
)

tokenizer.decode(trainer.train_dataset[5]["input_ids"])

space = tokenizer(" ", add_special_tokens = False).input_ids[0]


tokenizer.decode([space if x == -100 else x for x in trainer.train_dataset[5]["labels"]])

trainer_stats = trainer.train()

from unsloth.chat_templates import get_chat_template


tokenizer = get_chat_template(
  tokenizer,
  chat_template = "llama-3.1",
)
FastLanguageModel.for_inference(model) # Enable native 2x faster inference


messages = [
  {"role": "user", "content": "Continue the fibonnaci sequence: 1, 1, 2, 3, 5, 8,"},
]


inputs = tokenizer.apply_chat_template(
  messages,
  tokenize = True,
  add_generation_prompt = True, # Must add for generation
  return_tensors = "pt",
).to("cuda")


outputs = model.generate(input_ids = inputs, max_new_tokens = 64, use_cache = True,
  temperature = 1.5, min_p = 0.1)


tokenizer.batch_decode(outputs)

model.save_pretrained("lora_model") # Local saving
tokenizer.save_pretrained("lora_model")

if False:
    from unsloth import FastLanguageModel
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = "lora_model", # YOUR MODEL YOU USED FOR TRAINING
        max_seq_length = max_seq_length,
        dtype = dtype,
        load_in_4bit = load_in_4bit,
    )
    FastLanguageModel.for_inference(model) # Enable native 2x faster inference


messages = [
    {"role": "user", "content": "Describe a tall tower in the capital of France."},
]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize = True,
    add_generation_prompt = True, # Must add for generation
    return_tensors = "pt",
).to("cuda")


from transformers import TextStreamer
text_streamer = TextStreamer(tokenizer, skip_prompt = True)
_ = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128,
                   use_cache = True, temperature = 1.5, min_p = 0.1)