parser = argparse.ArgumentParser(description="Train Gemma model with LoRA")
parser.add_argument("--model_id", type=str, default="google/gemma-3-4b-it",
help="Base model ID (default: google/gemma-3-4b-it)")
parser.add_argument("--processor_id", type=str, default="google/gemma-3-4b-it",
help="Processor ID (default: google/gemma-3-4b-it)")
parser.add_argument("--train_jsonl", type=str, required=True,
help="Path to training JSONL file")
parser.add_argument("--output_dir", type=str, default="gemma-zipper-lora",
help="Output directory (default: gemma-zipper-lora)")
parser.add_argument("--hub_repo", type=str, default="ayushadarsh7/gemma3_lora",
help="HuggingFace Hub repository name (e.g., username/model-name)")
parser.add_argument("--num_epochs", type=int, default=3,
help="Number of training epochs (default: 3)")
parser.add_argument("--batch_size", type=int, default=1,
help="Batch size per device (default: 1)")
parser.add_argument("--gradient_accumulation_steps", type=int, default=4,
help="Gradient accumulation steps (default: 4)")
parser.add_argument("--learning_rate", type=float, default=2e-4,
help="Learning rate (default: 2e-4)")
parser.add_argument("--lora_r", type=int, default=16,
help="LoRA r parameter (default: 16)")
parser.add_argument("--lora_alpha", type=int, default=16,
help="LoRA alpha parameter (default: 16)")
parser.add_argument("--merge_and_save", action="store_true",
help="Merge LoRA adapter with base model and save")
