Xdimitris15gmailcom/

Di_mi_ai

# 1. Load Tokenizer and Model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# GPT-2 does not have a padding token by default, so we use the EOS token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 2. Prepare the Dataset
# This class reads your .txt file and chunks it into token blocks for the model
def load_dataset(path, tokenizer, block_size=128):
    return TextDataset(
        tokenizer=tokenizer,
        file_path=path,
        block_size=block_size,
    )

train_dataset = load_dataset(file_path, tokenizer)

# 3. Data Collator
# This prepares the batches and handles shifting labels for causal language modeling
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, 
    mlm=False # Causal LM, not Masked LM
)

# 4. Define Training Arguments
training_args = TrainingArguments(
    output_dir=output_dir,
    overwrite_output_dir=True,
    num_train_epochs=3,           # Number of times to go through your data
    per_device_train_batch_size=4, # Adjust based on your GPU/RAM
    save_steps=500,               # Save checkpoint every 500 steps
    save_total_limit=2,           # Only keep the 2 most recent checkpoints
    logging_steps=10,
    prediction_loss_only=True,
)

# 5. Initialize Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
)

# 6. Train and Save
print("Starting training...")
trainer.train()

trainer.save_model(output_dir)
tokenizer.save_pretrained(output_dir)
print(f"Training complete. Model saved to {output_dir}")

if not os.path.exists(data_file):
    with open(data_file, "w") as f:
        f.write("Artificial intelligence is a branch of computer science.\n" * 100)
    print(f"Created a dummy {data_file}. Replace this with your actual data!")

# Run the process
# NOTE: This requires 'pip install transformers torch accelerate'
finetune_on_text(data_file)

# Test the result
test_model("./fine_tuned_ai", prompt="AI is basically")