Create train.py

train.py

@@ -0,0 +1,101 @@
+print("Loading...")
+
+import torch
+
+torch.cuda.empty_cache()
+
+from datasets import load_dataset
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+
+MODEL_NAME = "Pin-25M"
+DATASET_ID = "starhopp3r/TinyChat"
+MAX_LENGTH = 256
+BATCH_SIZE = 32
+
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
+
+config = AutoConfig.from_pretrained(
+    "gpt2",
+    n_layer=12,
+    n_head=12,
+    n_embd=288,
+    n_inner=1152,
+    vocab_size=len(tokenizer),
+    bos_token_id=tokenizer.bos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+)
+model = AutoModelForCausalLM.from_config(config)
+
+print(f"Model parameters: {model.num_parameters() / 1e6:.2f}M")
+
+print("Loading dataset...")
+
+dataset = load_dataset(DATASET_ID, split="train")
+
+def tokenize_function(examples):
+    return tokenizer(examples["text"], truncation=True, max_length=MAX_LENGTH)
+
+tokenized_datasets = dataset.map(
+    tokenize_function, 
+    batched=True, 
+    remove_columns=dataset.column_names,
+    num_proc=4
+)
+
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+print("Setting up training arguments...")
+
+training_args = TrainingArguments(
+    output_dir="./" + MODEL_NAME + "_checkpoints",
+    num_train_epochs=1,
+    max_steps=1500,
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=2,
+    learning_rate=5e-4,
+    weight_decay=0.01,
+    logging_steps=100,
+    save_steps=2500,
+    fp16=True,
+    push_to_hub=False,
+    report_to="none",
+    warmup_steps=500,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets,
+    data_collator=data_collator,
+)
+
+print("Starting training...")
+trainer.train()
+
+trainer.save_model("./" + MODEL_NAME + "-Final")
+tokenizer.save_pretrained("./" + MODEL_NAME + "-Final")
+
+def chat(prompt):
+    formatted_prompt = f"[INST] {prompt} [/INST]"
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
+    model.to("cuda")
+    
+    outputs = model.generate(
+        **inputs, 
+        max_new_tokens=50, 
+        temperature=0.7, 
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+print("\n--- Test Chat ---")
+print(chat("Hello, how are you today?"))