Upload train_human.py

train_human.py

@@ -0,0 +1,146 @@
+# %%
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from datasets import load_dataset, Dataset
+
+from trl import DPOTrainer, DPOConfig
+from peft import LoraConfig
+from peft import prepare_model_for_kbit_training
+import torch
+
+import pandas as pd
+
+# %%
+dataset = load_dataset("Undi95/Weyaxi-humanish-dpo-project-noemoji")["train"]
+
+model_name = "Undi95/Meta-Llama-3.1-8B-Claude-bf16"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.padding_side = "right"
+
+tokenizer.pad_token = tokenizer.eos_token
+
+# %%
+tokenizer.chat_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}"
+
+
+# %%
+dataset2 = load_dataset("ResplendentAI/NSFW_RP_Format_DPO")['train']
+
+# %%
+dataset = dataset.to_pandas(
+)
+dataset2 = dataset2.to_pandas()
+
+dataset = Dataset.from_pandas(pd.concat([dataset.sample(400), dataset2]).sample(frac=1))
+
+# %%
+def template_prompt(system, prompt):
+    if system is None:
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+    else:
+        messages = [
+            {
+                "role": "system",
+                "content": system,
+            },
+            {"role": "user", "content": prompt},
+        ]
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=False
+    )
+    return prompt
+
+
+def template_answer(answer):
+    messages = [
+        {
+            "role": "assistant",
+            "content": answer,
+        },
+    ]
+    answer = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=False
+    )
+    return answer
+
+# %%
+# create new columns
+dataset = dataset.map(
+    lambda x: {
+        "prompt": template_prompt(None, x["prompt"]).replace("<|start_header_id|>assistant<|end_header_id|>\n\n", "")
+    },  # change this according to the dataset!!!
+)
+
+# %%
+dataset = dataset.map(
+    lambda x: {"chosen": template_answer(x["chosen"]).replace('<|begin_of_text|>', '').replace('><|start_header_id|>assistant<|end_header_id|>\n\n', '>')},
+)
+dataset = dataset.map(
+    lambda x: {"rejected": template_answer(x["rejected"]).replace('<|begin_of_text|>', '').replace('><|start_header_id|>assistant<|end_header_id|>\n\n', '>')},
+)
+
+# %%
+dataset[0]
+
+# %%
+# LoRA configuration
+peft_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=[
+        "k_proj",
+        "gate_proj",
+        "v_proj",
+        "up_proj",
+        "q_proj",
+        "o_proj",
+        "down_proj",
+    ],
+)
+
+# Model to fine-tune
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    load_in_4bit=True,
+    device_map="auto",
+)
+model.config.use_cache = False
+
+
+model.gradient_checkpointing_enable()
+model = prepare_model_for_kbit_training(model)
+
+# %%
+output_name = f"checkpoints/exp_human_{model_name}"
+
+training_args = DPOConfig(
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=4,
+    num_train_epochs=1,
+    gradient_checkpointing=True,
+    output_dir=output_name,
+    logging_steps=1,
+    max_steps=50
+)
+
+trainer = DPOTrainer(
+    model,
+    ref_model=None,
+    train_dataset=dataset,
+    tokenizer=tokenizer,
+    args=training_args,
+    peft_config=peft_config,
+)
+
+trainer.train()
+
+trainer.save_model(output_name)
+
+