%%capture
# In Colab, we skip dependency installation to avoid conflicts with preinstalled packages.
# On local machines, we include dependencies for completeness.

import os
if "COLAB_" not in "".join(os.environ.keys()):
    !pip install unsloth vllm==0.8.2
else:

    !pip install --no-deps unsloth vllm==0.8.2

# Get https://github.com/meta-llama/synthetic-data-kit
!pip install synthetic-data-kit

%%capture
import os
if "COLAB_" in "".join(os.environ.keys()):

    import sys, re, requests; modules = list(sys.modules.keys())
    for x in modules: sys.modules.pop(x) if "PIL" in x or "google" in x else None
    !pip install --no-deps bitsandbytes accelerate xformers==0.0.29.post3 peft "trl==0.15.2" triton cut_cross_entropy unsloth_zoo
    !pip install sentencepiece protobuf datasets huggingface_hub[hf_xet] hf_transfer

    # vLLM requirements - vLLM breaks Colab due to reinstalling numpy
    f = requests.get("https://raw.githubusercontent.com/vllm-project/vllm/refs/heads/main/requirements/common.txt").content
    with open("vllm_requirements.txt", "wb") as file:
        file.write(re.sub(rb"(transformers|numpy|xformers|importlib_metadata)[^\n]{0,}\n", b"", f))
    !pip install -r vllm_requirements.txt

>>> from unsloth.dataprep import SyntheticDataKit

>>> generator = SyntheticDataKit.from_pretrained(
...     model_name="unsloth/Llama-3.2-3B-Instruct",
...     max_seq_length=2048,
... )

generator.prepare_qa_generation(
    output_folder="data",  # Output location of synthetic data
    temperature=0.7,  # Higher temp makes more diverse data
    top_p=0.95,
    overlap=64,  # Overlap portion during chunking
    max_generation_tokens=512,  # Can increase for longer QA pairs
)

>>> !synthetic-data-kit system-check

>>> # Make sure synthetic_data_kit_config.yaml points to the 'data_docs' folder
>>> !synthetic-data-kit -c synthetic_data_kit_config.yaml ingest /content/langchain-ai-langchain.txt

>>> filenames = generator.chunk_data("data/output/langchain-ai-langchain.txt")
>>> print(f"Created {len(filenames)} chunks.")

import time
# Process 2 chunks for now -> can increase but slower!
for filename in filenames[:2]:
    !synthetic-data-kit \
        -c synthetic_data_kit_config.yaml \
        create {filename} \
        --num-pairs 25 \
        --type "qa"
    time.sleep(2) # Sleep some time to leave some room for processing

>>> qa_pairs_filenames = [
...     f"data/generated/langchain-ai-langchain_{i}_qa_pairs.json"
...     for i in range(len(filenames[:2]))
... ]
>>> for filename in qa_pairs_filenames:
...     !synthetic-data-kit \
...         -c synthetic_data_kit_config.yaml \
...         save-as {filename} -f ft

>>> generator.cleanup()

from datasets import Dataset
import pandas as pd

final_filenames = [f"data/final/langchain-ai-langchain_{i}_qa_pairs_ft.json" for i in range(len(filenames[:2]))]
conversations = pd.concat([pd.read_json(name) for name in final_filenames]).reset_index(drop=True)

dataset = Dataset.from_pandas(conversations)

dataset[0]

dataset[-1]

# !zip -r data.zip /content/
# !unzip data.zip

# import os
# import pandas as pd
# from datasets import Dataset

# # Path to your folder containing JSON files
# folder_path = 'content/data/final/'

# # List all .json files in the folder
# final_filenames = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith('.json')]

# # Read and combine the JSON files
# conversations = pd.concat([
#     pd.read_json(name) for name in final_filenames
# ]).reset_index(drop=True)

# # Convert to Hugging Face Dataset
# dataset = Dataset.from_pandas(conversations)

>>> from unsloth import FastLanguageModel
>>> import torch


>>> model, tokenizer = FastLanguageModel.from_pretrained(
...     model_name="unsloth/Llama-3.2-3B-Instruct",
...     max_seq_length=1024,  # Choose any for long context!
...     load_in_4bit=True,  # 4 bit quantization to reduce memory
...     load_in_8bit=False,  # [NEW!] A bit more accurate, uses 2x memory
...     full_finetuning=False,  # [NEW!] We have full finetuning now!
... )

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
)

<|begin_of_text|><|start_header_id|>system<|end_header_id|>

Cutting Knowledge Date: December 2023
Today Date: 01 May 2025

You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>

What is 1+1?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

2<|eot_id|>

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


pass

# Get our previous dataset and format it:
dataset = dataset.map(
    formatting_prompts_func,
    batched=True,
)

dataset[0]

from trl import SFTTrainer, SFTConfig

trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    train_dataset=dataset,
    eval_dataset=None,  # Can set up evaluation!
    args=SFTConfig(
        dataset_text_field="text",
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,  # Use GA to mimic batch size!
        warmup_steps=5,
        max_steps=60,
        learning_rate=2e-4,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        report_to="none",  # Use this for WandB etc
    ),
)

trainer_stats = trainer.train()

>>> messages = [
...     {"role": "user", "content": "What is the standard interface for binding tools to models?"},
... ]
>>> 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=256, temperature=0.1)