Upload scripts/train_tokenizer.py with huggingface_hub

scripts/train_tokenizer.py

@@ -3,7 +3,7 @@
 Command-line interface for training BBPE tokenizers.
 
 Usage:
-    python train_tokenizer.py --data_dir ./data --vocab_size 30000 --model_name my_tokenizer
+    python train_tokenizer.py --data_dir ./data --vocab_size 30000 --model_name EthioBBPE
 """
 
 import argparse
@@ -13,7 +13,7 @@ from pathlib import Path
 # Add parent directory to path for imports
 sys.path.insert(0, str(Path(__file__).parent))
 
-from bbpe_trainer import BBPETrainer, BBPEConfig
+from bbpe_trainer import EthioBBPETrainer, BBPEConfig
 
 
 def parse_args():
@@ -100,7 +100,7 @@ def parse_args():
     parser.add_argument(
         "--model_name",
         type=str,
-        default="bbpe_tokenizer",
+        default="EthioBBPE",
         help="Name for the saved tokenizer model",
     )
     
@@ -119,6 +119,88 @@ def parse_args():
         help="Path to save the configuration JSON file",
     )
     
+    # Advanced production features
+    parser.add_argument(
+        "--use_checkpoint",
+        action="store_true",
+        default=True,
+        help="Enable checkpointing during training",
+    )
+    
+    parser.add_argument(
+        "--no_checkpoint",
+        action="store_false",
+        dest="use_checkpoint",
+        help="Disable checkpointing",
+    )
+    
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=str,
+        default="./models/checkpoints",
+        help="Directory to save checkpoints",
+    )
+    
+    parser.add_argument(
+        "--max_checkpoints",
+        type=int,
+        default=5,
+        help="Maximum number of checkpoints to keep (0 = unlimited)",
+    )
+    
+    parser.add_argument(
+        "--save_compressed",
+        action="store_true",
+        default=True,
+        help="Save tokenizer files in compressed format (.gz, .bz2, .xz)",
+    )
+    
+    parser.add_argument(
+        "--no_compression",
+        action="store_false",
+        dest="save_compressed",
+        help="Disable compression",
+    )
+    
+    parser.add_argument(
+        "--compression_format",
+        type=str,
+        choices=['gzip', 'bz2', 'lzma'],
+        default='gzip',
+        help="Compression format to use",
+    )
+    
+    parser.add_argument(
+        "--compression_level",
+        type=int,
+        choices=range(1, 10),
+        default=9,
+        help="Compression level (1-9, higher = better compression but slower)",
+    )
+    
+    parser.add_argument(
+        "--enable_quantization",
+        action="store_true",
+        default=False,
+        help="Enable model quantization for deployment",
+    )
+    
+    parser.add_argument(
+        "--quantization_bits",
+        type=int,
+        choices=[4, 8],
+        default=8,
+        help="Quantization bits (4 or 8)",
+    )
+    
+    parser.add_argument(
+        "--export_formats",
+        type=str,
+        nargs='+',
+        default=['tokenizer.json', 'vocab_compressed', 'hf_export'],
+        help="Export formats to generate",
+    )
+    
     return parser.parse_args()
 
 
@@ -137,11 +219,18 @@ def main():
             min_frequency=args.min_frequency,
             special_tokens=args.special_tokens,
             lowercase=args.lowercase,
-            add_prefix_space=not args.no_prefix_space,
             show_progress=args.show_progress,
             data_dir=args.data_dir,
             model_save_dir=args.model_save_dir,
             model_name=args.model_name,
+            use_checkpoint=args.use_checkpoint,
+            checkpoint_dir=args.checkpoint_dir,
+            save_compressed=args.save_compressed,
+            compression_format=args.compression_format,
+            compression_level=args.compression_level,
+            enable_quantization=args.enable_quantization,
+            quantization_bits=args.quantization_bits,
+            max_checkpoints=args.max_checkpoints,
         )
     
     # Save config if requested
@@ -150,7 +239,7 @@ def main():
         print(f"Configuration saved to {args.save_config}")
     
     # Initialize trainer
-    trainer = BBPETrainer(config)
+    trainer = EthioBBPETrainer(config)
     
     # Get training files
     if args.files:
@@ -171,7 +260,21 @@ def main():
         sys.exit(1)
     
     # Save the tokenizer
-    save_path = trainer.save()
+    save_path = trainer.save(export_formats=args.export_formats)
+    
+    # List checkpoints
+    print("\n" + "="*60)
+    print("CHECKPOINTS")
+    print("="*60)
+    checkpoints = trainer.list_checkpoints()
+    if checkpoints:
+        for ckpt in checkpoints:
+            status = "✓" if trainer.validate_checkpoint(ckpt['path']) else "✗"
+            print(f"{status} {ckpt['name']} ({ckpt['size_kb']:.1f} KB)")
+            if 'vocab_size' in ckpt:
+                print(f"   Vocab size: {ckpt['vocab_size']}, Final: {ckpt.get('is_final', False)}")
+    else:
+        print("No checkpoints found.")
     
     # Test the tokenizer
     print("\n" + "="*60)
@@ -180,7 +283,7 @@ def main():
     
     test_texts = [
         "Hello, world!",
-        "This is a test of the BBPE tokenizer.",
+        "This is a test of the EthioBBPE tokenizer.",
         "Special characters: @#$%^&*()",
         "Numbers: 12345 and words mixed together.",
     ]
@@ -198,6 +301,25 @@ def main():
     print("\n" + "="*60)
     print(f"Tokenizer training complete!")
     print(f"Model saved to: {save_path}")
+    if args.save_compressed:
+        print(f"Compressed files saved (format: {args.compression_format}, level: {args.compression_level})")
+    if args.enable_quantization:
+        print(f"Quantized model saved ({args.quantization_bits}-bit)")
+    if args.use_checkpoint:
+        print(f"Checkpoints saved to: {args.checkpoint_dir} (max: {args.max_checkpoints})")
+    
+    # Print training metrics
+    print("\n" + "="*60)
+    print("TRAINING METRICS")
+    print("="*60)
+    metrics = trainer.get_training_metrics()
+    if 'initial' in metrics:
+        print(f"Files processed: {metrics['initial'].get('num_files', 'N/A')}")
+        print(f"Total data size: {metrics['initial'].get('total_bytes', 0) / 1024:.2f} KB")
+    if 'final' in metrics:
+        print(f"Final vocab size: {metrics['final'].get('vocab_size', 'N/A')}")
+        print(f"Training duration: {metrics['final'].get('training_duration_sec', 0):.2f} seconds")
+    
     print("="*60)