from accelerate import PartialState  # type: ignore
from datasets import load_dataset  # type: ignore
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer  # type: ignore

device_string = PartialState().process_index
# MODEL = "meta-llama/Llama-3.3-70B-Instruct"
MODEL = "Qwen/Qwen-8b"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="auto")
output_dir = "../weights_5/"
dataset = load_dataset("json", data_files="cleaned_dataset.jsonl", split="train")
# Format messages into a flat text string for models like OPT
# chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"

# model, tokenizer = setup_chat_format(
#     model, tokenizer, "chatml")

training_args = SFTConfig(
    output_dir=output_dir,
    max_seq_length=512,
    gradient_checkpointing_kwargs={"use_reentrant": False},
)

trainer = SFTTrainer(
    model=model, args=training_args, train_dataset=dataset, processing_class=tokenizer
)

trainer.train(resume_from_checkpoint=False)

trainer.save_model(output_dir)
tokenizer.save_pretrained(output_dir)
