init: add LoRA training infrastructure and 20k dataset
This commit is contained in:
Christian Medina
103
training/scripts/inference.py
Executable file
103
training/scripts/inference.py
Executable file
@@ -0,0 +1,103 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Inference script for trained LoRA adapter.
|
||||
|
||||
Loads the trained model and generates Cyron summaries.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def load_model(model_path):
|
||||
"""Load trained LoRA model."""
|
||||
from peft import PeftModel
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
print(f"Loading base model from {model_path}...")
|
||||
base_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
device_map="auto",
|
||||
torch_dtype="auto",
|
||||
)
|
||||
|
||||
print(f"Loading LoRA weights from {model_path}...")
|
||||
model = PeftModel.from_pretrained(base_model, model_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def generate_summary(model, tokenizer, task, files_changed=None, tests_run=False, test_count=0, commit=False, push=False, errors_seen=False):
|
||||
"""Generate a Cyron summary for a task."""
|
||||
|
||||
# Build prompt
|
||||
prompt = f"Task: {task}"
|
||||
if files_changed:
|
||||
prompt += f"\nFiles: {', '.join(files_changed[:3])}"
|
||||
|
||||
if tests_run:
|
||||
prompt += f"\nTests run: {test_count}"
|
||||
if commit:
|
||||
prompt += "\nCommit: yes"
|
||||
if push:
|
||||
prompt += "\nPush: yes"
|
||||
if errors_seen:
|
||||
prompt += "\nErrors seen: yes"
|
||||
|
||||
prompt += "\n\nGenerate a Cyron summary:"
|
||||
|
||||
# Tokenize
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
||||
|
||||
# Generate
|
||||
outputs = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=500,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
repetition_penalty=1.1,
|
||||
)
|
||||
|
||||
# Decode
|
||||
generated = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||||
|
||||
return generated
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate Cyron summaries with LoRA")
|
||||
parser.add_argument("--model", type=str, default="output/llama2-7b-lora",
|
||||
help="Path to trained model")
|
||||
parser.add_argument("--task", type=str, required=True, help="Task description")
|
||||
parser.add_argument("--files", type=str, nargs="*", default=[], help="Files changed")
|
||||
parser.add_argument("--tests-run", action="store_true", help="Tests were run")
|
||||
parser.add_argument("--test-count", type=int, default=0, help="Number of tests")
|
||||
parser.add_argument("--commit", action="store_true", help="Commit was made")
|
||||
parser.add_argument("--push", action="store_true", help="Code was pushed")
|
||||
parser.add_argument("--errors", action="store_true", help="Errors were seen")
|
||||
args = parser.parse_args()
|
||||
|
||||
model, tokenizer = load_model(args.model)
|
||||
|
||||
summary = generate_summary(
|
||||
model,
|
||||
tokenizer,
|
||||
task=args.task,
|
||||
files_changed=args.files,
|
||||
tests_run=args.tests_run,
|
||||
test_count=args.test_count,
|
||||
commit=args.commit,
|
||||
push=args.push,
|
||||
errors_seen=args.errors,
|
||||
)
|
||||
|
||||
print("\nGenerated Summary:")
|
||||
print("=" * 70)
|
||||
print(summary)
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
95
training/scripts/prepare_dataset.py
Executable file
95
training/scripts/prepare_dataset.py
Executable file
@@ -0,0 +1,95 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Prepare Cyron summary dataset for LoRA training.
|
||||
|
||||
Reads combined_20k.jsonl and formats it for Hugging Face TRL training.
|
||||
"""
|
||||
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def prepare_dataset(input_file, output_file, test_size=0.05):
|
||||
"""Prepare dataset for training."""
|
||||
|
||||
with open(input_file) as f:
|
||||
examples = [json.loads(line) for line in f]
|
||||
|
||||
print(f"Loaded {len(examples)} examples")
|
||||
|
||||
# Format for training
|
||||
formatted = []
|
||||
for ex in examples:
|
||||
# Create instruction-input-output format
|
||||
instruction = f"Task: {ex['task']}"
|
||||
if ex['files_changed']:
|
||||
instruction += f"\nFiles: {', '.join(ex['files_changed'][:3])}"
|
||||
|
||||
input_text = ""
|
||||
if ex['tests_run']:
|
||||
input_text += f"Tests run: {ex['test_count']}"
|
||||
if ex['commit']:
|
||||
input_text += f"\nCommit: {ex['git']['commit'] if ex.get('git') else 'yes'}"
|
||||
if ex['push']:
|
||||
input_text += "\nPush: yes"
|
||||
|
||||
output_text = ex['output']
|
||||
|
||||
# Create conversation format
|
||||
conversation = {
|
||||
"conversations": [
|
||||
{
|
||||
"from": "human",
|
||||
"value": f"Generate a Cyron summary for this task:\n\n{instruction}\n\n{input_text}"
|
||||
},
|
||||
{
|
||||
"from": "gpt",
|
||||
"value": output_text
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
formatted.append(conversation)
|
||||
|
||||
# Split into train/test
|
||||
import random
|
||||
random.seed(42)
|
||||
random.shuffle(formatted)
|
||||
|
||||
split_point = int(len(formatted) * (1 - test_size))
|
||||
train_data = formatted[:split_point]
|
||||
test_data = formatted[split_point:]
|
||||
|
||||
# Save
|
||||
with open(output_file.parent / "train.jsonl", "w") as f:
|
||||
for item in train_data:
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
with open(output_file.parent / "test.jsonl", "w") as f:
|
||||
for item in test_data:
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
print(f"Train: {len(train_data)} examples")
|
||||
print(f"Test: {len(test_data)} examples")
|
||||
print(f"Saved to {output_file.parent}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Prepare LoRA training dataset")
|
||||
parser.add_argument("--input", type=str, default="../combined_20k.jsonl",
|
||||
help="Input combined dataset")
|
||||
parser.add_argument("--output", type=str, default="data",
|
||||
help="Output directory")
|
||||
parser.add_argument("--test-size", type=float, default=0.05,
|
||||
help="Test set percentage")
|
||||
args = parser.parse_args()
|
||||
|
||||
output_dir = Path(args.output)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
prepare_dataset(args.input, output_dir, args.test_size)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
127
training/scripts/train.py
Executable file
127
training/scripts/train.py
Executable file
@@ -0,0 +1,127 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Train LoRA adapter on Cyron summary dataset.
|
||||
|
||||
Uses Hugging Face TRL for SFT training with QLoRA.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def train(config_path):
|
||||
"""Train LoRA adapter using TRL."""
|
||||
|
||||
with open(config_path) as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
BitsAndBytesConfig,
|
||||
TrainingArguments,
|
||||
)
|
||||
from peft import (
|
||||
prepare_model_for_kbit_training,
|
||||
LoraConfig,
|
||||
get_peft_model,
|
||||
)
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
print(f"Loading model: {config['base_model']}")
|
||||
|
||||
# Load model with quantization
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=config.get("load_in_4bit", True),
|
||||
bnb_4bit_compute_dtype=config.get("bnb_4bit_compute_dtype", "bfloat16"),
|
||||
bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"),
|
||||
use_nested_quant=config.get("use_nested_quant", False),
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
config["base_model"],
|
||||
quantization_config=bnb_config,
|
||||
device_map="auto",
|
||||
)
|
||||
model = prepare_model_for_kbit_training(model)
|
||||
|
||||
# Add LoRA
|
||||
lora_config = LoraConfig(
|
||||
r=config["lora_r"],
|
||||
lora_alpha=config["lora_alpha"],
|
||||
lora_dropout=config["lora_dropout"],
|
||||
target_modules=config["target_modules"],
|
||||
task_type=config.get("lora_task_type", "CAUSAL_LM"),
|
||||
)
|
||||
model = get_peft_model(model, lora_config)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
# Load tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(config["base_model"])
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.padding_side = "right"
|
||||
|
||||
# Load dataset
|
||||
from datasets import load_dataset
|
||||
|
||||
dataset = load_dataset(
|
||||
"json",
|
||||
data_files={
|
||||
"train": config["dataset"][0]["path"].replace("../", ""),
|
||||
"test": config["dataset"][0]["path"].replace("../", "").replace("combined_20k.jsonl", "test.jsonl"),
|
||||
},
|
||||
)
|
||||
|
||||
# Training arguments
|
||||
training_args = TrainingArguments(
|
||||
output_dir=config["train_params"]["output_dir"],
|
||||
num_train_epochs=config["train_params"]["num_train_epochs"],
|
||||
per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
|
||||
gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
|
||||
learning_rate=config["train_params"]["learning_rate"],
|
||||
lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
|
||||
weight_decay=config["train_params"]["weight_decay"],
|
||||
warmup_ratio=config["train_params"]["warmup_ratio"],
|
||||
max_seq_length=config["train_params"]["max_seq_length"],
|
||||
logging_steps=config["train_params"]["logging_steps"],
|
||||
save_steps=config["train_params"]["save_steps"],
|
||||
save_total_limit=config["train_params"]["save_total_limit"],
|
||||
evaluation_strategy=config.get("eval_strategy", "steps"),
|
||||
eval_steps=config.get("eval_steps", 100),
|
||||
mixed_precision=config.get("mixed_precision", "bf16"),
|
||||
gradient_checkpointing=config.get("gradient_checkpointing", True),
|
||||
)
|
||||
|
||||
# SFT Trainer
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset["train"],
|
||||
eval_dataset=dataset["test"],
|
||||
args=training_args,
|
||||
max_seq_length=config["train_params"]["max_seq_length"],
|
||||
)
|
||||
|
||||
# Train
|
||||
print("Starting training...")
|
||||
trainer.train()
|
||||
|
||||
# Save
|
||||
trainer.save_model(config["train_params"]["output_dir"])
|
||||
tokenizer.save_pretrained(config["train_params"]["output_dir"])
|
||||
|
||||
print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Train LoRA adapter")
|
||||
parser.add_argument("--config", type=str, default="configs/llama2-7b-lora.yaml",
|
||||
help="Training configuration file")
|
||||
args = parser.parse_args()
|
||||
|
||||
train(args.config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user