import torch
import datasets
from transformers import TrainingArguments, AutoConfig, AutoTokenizer, AutoModelForCausalLM
import trl
from transformers import BitsAndBytesConfig

train_dataset = datasets.load_dataset('HuggingFaceTB/cosmopedia-20k', split='train')

args = TrainingArguments(
    output_dir="./test-sft",
    max_steps=20000,
    per_device_train_batch_size=1,
    optim="adafactor", report_to="none",
)

model_id = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
nf4_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config,device_map="auto")
print(model)

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64, target_modules=["q_proj", "v_proj", "k_proj", "dense", "lm_head", "fc1", "fc2"],
    bias="none", 
    task_type="CAUSAL_LM",
)
model.add_adapter(peft_config)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=1024
)

trainer.train()

trainer.model.save_pretrained("sft", dtype=torch.bfloat16)
trainer.tokenizer.save_pretrained("sft")