130 lines
4.8 KiB
Python
130 lines
4.8 KiB
Python
import numpy as np
|
|
import torch
|
|
from datasets import load_dataset
|
|
from transformers import (
|
|
AutoTokenizer,
|
|
AutoModelForCausalLM,
|
|
TrainingArguments,
|
|
Trainer,
|
|
DataCollatorForLanguageModeling,
|
|
)
|
|
|
|
# ─── Configuration ───────────────────────────────────────────────────────────
|
|
MODEL_NAME = "zxc4wewewe/blackthinking" # lightweight model suitable for CPU
|
|
MAX_LENGTH = 512 # max token length per example
|
|
OUTPUT_DIR = "./results"
|
|
NUM_EPOCHS = 3
|
|
BATCH_SIZE = 2 # small batch for CPU training
|
|
LEARNING_RATE = 5e-5
|
|
LOGGING_STEPS = 50
|
|
|
|
# ─── 1. Load dataset from Hugging Face Hub ───────────────────────────────────
|
|
dataset = load_dataset("zxc4wewewe/offsec")
|
|
print(f"Train: {len(dataset['train'])} examples | Test: {len(dataset['test'])} examples")
|
|
print(f"Columns: {dataset['train'].column_names}")
|
|
|
|
|
|
# ─── 2. Format & tokenize ────────────────────────────────────────────────────
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
|
|
|
# GPT-2 has no pad token by default — use eos_token
|
|
if tokenizer.pad_token is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
|
|
def format_and_tokenize(examples):
|
|
"""Combine prompt + response into a single text and tokenize."""
|
|
texts = [
|
|
f"{prompt}{response}{tokenizer.eos_token}"
|
|
for prompt, response in zip(examples["prompt"], examples["response"])
|
|
]
|
|
tokenized = tokenizer(
|
|
texts,
|
|
truncation=True,
|
|
max_length=MAX_LENGTH,
|
|
padding="max_length",
|
|
)
|
|
# For causal LM, labels = input_ids (the model learns to predict next token)
|
|
tokenized["labels"] = tokenized["input_ids"].copy()
|
|
return tokenized
|
|
|
|
|
|
tokenized_dataset = dataset.map(
|
|
format_and_tokenize,
|
|
batched=True,
|
|
remove_columns=dataset["train"].column_names,
|
|
desc="Tokenizing",
|
|
)
|
|
|
|
print(f"Tokenized train: {len(tokenized_dataset['train'])} examples")
|
|
|
|
|
|
# ─── 3. Model ────────────────────────────────────────────────────────────────
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
|
|
model.resize_token_embeddings(len(tokenizer))
|
|
|
|
data_collator = DataCollatorForLanguageModeling(
|
|
tokenizer=tokenizer,
|
|
mlm=False, # causal LM, not masked LM
|
|
)
|
|
|
|
|
|
# ─── 4. Training ─────────────────────────────────────────────────────────────
|
|
training_args = TrainingArguments(
|
|
output_dir=OUTPUT_DIR,
|
|
overwrite_output_dir=True,
|
|
num_train_epochs=NUM_EPOCHS,
|
|
per_device_train_batch_size=BATCH_SIZE,
|
|
per_device_eval_batch_size=BATCH_SIZE,
|
|
eval_strategy="epoch",
|
|
save_strategy="epoch",
|
|
learning_rate=LEARNING_RATE,
|
|
weight_decay=0.01,
|
|
logging_dir="./logs",
|
|
logging_steps=LOGGING_STEPS,
|
|
load_best_model_at_end=True,
|
|
save_total_limit=2,
|
|
fp16=False, # CPU-only
|
|
report_to="none",
|
|
)
|
|
|
|
trainer = Trainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=tokenized_dataset["train"],
|
|
eval_dataset=tokenized_dataset["test"],
|
|
data_collator=data_collator,
|
|
)
|
|
|
|
print("Starting training...")
|
|
trainer.train()
|
|
|
|
# Save final model
|
|
trainer.save_model(f"{OUTPUT_DIR}/final_model")
|
|
tokenizer.save_pretrained(f"{OUTPUT_DIR}/final_model")
|
|
print(f"Model saved to {OUTPUT_DIR}/final_model")
|
|
|
|
|
|
# ─── 5. Inference ────────────────────────────────────────────────────────────
|
|
def generate_response(prompt_text, max_new_tokens=256):
|
|
"""Generate a response given a prompt."""
|
|
inputs = tokenizer(prompt_text, return_tensors="pt")
|
|
with torch.no_grad():
|
|
output_ids = model.generate(
|
|
**inputs,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=True,
|
|
temperature=0.7,
|
|
top_p=0.9,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
# Decode only the generated part (skip the prompt tokens)
|
|
generated = output_ids[0][inputs["input_ids"].shape[1]:]
|
|
return tokenizer.decode(generated, skip_special_tokens=True)
|
|
|
|
|
|
# Example usage (uncomment to test after training):
|
|
sample_prompt = dataset["test"][0]["prompt"]
|
|
print("Prompt:", sample_prompt[:200], "...")
|
|
print("Generated:", generate_response(sample_prompt))
|