初始化项目,由ModelHub XC社区提供模型
Model: dejanseo/reverse-prompter Source: Original Platform
This commit is contained in:
68
assets/train-270m-ft.py
Normal file
68
assets/train-270m-ft.py
Normal file
@@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Full fine-tune gemma-3-270m to reconstruct prompts from model outputs."""
|
||||
|
||||
from datasets import load_from_disk
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
|
||||
|
||||
# --- Configuration ---
|
||||
MODEL_NAME = "google/gemma-3-270m"
|
||||
DATASET_PATH = "tokenized-dataset-plain"
|
||||
OUTPUT_DIR = "checkpoints-270m-ft"
|
||||
MAX_SEQ_LENGTH = 2048
|
||||
|
||||
# --- Training ---
|
||||
BATCH_SIZE = 2
|
||||
GRAD_ACCUM = 8
|
||||
LEARNING_RATE = 5e-5
|
||||
NUM_EPOCHS = 1
|
||||
WARMUP_STEPS = 100
|
||||
LOGGING_STEPS = 1
|
||||
SAVE_STEPS = 100
|
||||
|
||||
|
||||
def main():
|
||||
print("Loading pre-tokenized dataset...")
|
||||
dataset = load_from_disk(DATASET_PATH)
|
||||
print(f"Training examples: {len(dataset)}")
|
||||
|
||||
print(f"Loading model: {MODEL_NAME}")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_NAME,
|
||||
torch_dtype="bfloat16",
|
||||
device_map="auto",
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
|
||||
trainer = Trainer(
|
||||
model=model,
|
||||
processing_class=tokenizer,
|
||||
train_dataset=dataset,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True),
|
||||
args=TrainingArguments(
|
||||
output_dir=OUTPUT_DIR,
|
||||
per_device_train_batch_size=BATCH_SIZE,
|
||||
gradient_accumulation_steps=GRAD_ACCUM,
|
||||
learning_rate=LEARNING_RATE,
|
||||
num_train_epochs=NUM_EPOCHS,
|
||||
warmup_steps=WARMUP_STEPS,
|
||||
logging_steps=LOGGING_STEPS,
|
||||
save_steps=SAVE_STEPS,
|
||||
bf16=True,
|
||||
seed=42,
|
||||
report_to="wandb",
|
||||
logging_strategy="steps",
|
||||
gradient_checkpointing=True,
|
||||
),
|
||||
)
|
||||
|
||||
print("Training...")
|
||||
trainer.train()
|
||||
|
||||
print("Saving model...")
|
||||
trainer.save_model(f"{OUTPUT_DIR}/final")
|
||||
tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
3
assets/train-loss.png
Normal file
3
assets/train-loss.png
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:22e27d2504371f7a5c14d1cf816fe3a84e2fb548308a517496d36db303f44165
|
||||
size 116660
|
||||
Reference in New Issue
Block a user