Buckets:
| 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)) | |
Xet Storage Details
- Size:
- 4.92 kB
- Xet hash:
- 0aeb75ae03c306ae329e19a73b6050d40d6a6895bda1b8e69aa2f796db59ac1d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.