Update arguments

hf-quantize.sh

@@ -4,14 +4,14 @@
 #
 # Usage with long options:
 #   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M
-#   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix true
-#   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix true --output-filename Llama-2-7b-Q4_K_M.gguf
-#   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix true --output-filename Llama-2-7b-Q4_K_M.gguf --split-model true --split-max-tensors 256 --split-max-size 4G
+#   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix
+#   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix --output-filename Llama-2-7b-Q4_K_M.gguf
+#   ./hf-quantize.sh --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix --output-filename Llama-2-7b-Q4_K_M.gguf --split-model --split-max-tensors 256 --split-max-size 4G
 #
 #   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M
-#   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M -i true
-#   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M -i true -o Llama-2-7b-Q4_K_M.gguf
-#   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M -i true -o Llama-2-7b-Q4_K_M.gguf -s true -t 256 -z 4G
+#   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M -imatrix
+#   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M -imatrix -o Llama-2-7b-Q4_K_M.gguf
+#   ./hf-quantize.sh -m meta-llama/Llama-2-7b -q Q4_K_M -imatrix -o Llama-2-7b-Q4_K_M.gguf -split --split-max-tensors 256 --split-max-size 4G
 #
 
 # --- Configuration ---
@@ -43,6 +43,15 @@ SPLIT_MAX_TENSORS="256"
 # Optional: Max size per shard if splitting (e.g., 2G) - overrides SPLIT_MAX_TENSORS if set
 SPLIT_MAX_SIZE=""
 
+# Optional: Quant embeddings tensor
+TOKEN_EMBEDDING_TYPE=""
+
+# Optional: Leave output tensor
+LEAVE_OUTPUT_TENSOR="false"
+
+# Optional: Output Quantization Method
+OUTPUT_TENSOR_TYPE=""
+
 # --- Parse Named Arguments ---
 while [[ $# -gt 0 ]]; do
     case $1 in
@@ -54,39 +63,51 @@ while [[ $# -gt 0 ]]; do
             QUANT_METHOD="$2"
             shift 2
             ;;
-        -i|--use-imatrix)
-            USE_IMATRIX="$2"
-            shift 2
+        -imatrix|--use-imatrix)
+            USE_IMATRIX="true"
+            shift 1
             ;;
         -o|--output-filename)
             OUTPUT_FILENAME="$2"
             shift 2
             ;;
-        -s|--split-model)
-            SPLIT_MODEL="$2"
-            shift 2
+        -split|--split-model)
+            SPLIT_MODEL="true"
+            shift 1
             ;;
-        -t|--split-max-tensors)
+        --split-max-tensors)
             SPLIT_MAX_TENSORS="$2"
             shift 2
             ;;
-        -z|--split-max-size)
+        --split-max-size)
             SPLIT_MAX_SIZE="$2"
             shift 2
             ;;
+        --token-embedding-type)
+            TOKEN_EMBEDDING_TYPE="$2"
+            shift 2
+            ;;
+        --leave-output-tensor)
+            LEAVE_OUTPUT_TENSOR="true"
+            shift 1
+            ;;
+        --output-tensor-type)
+            OUTPUT_TENSOR_TYPE="$2"
+            shift 2
+            ;;
         -h|--help)
             echo "Usage:"
             echo "  Long options:"
-            echo "    $0 --model <MODEL_ID> --quant-method <QUANT_METHOD> [--use-imatrix <true|false>] [--output-filename <FILENAME>] [--split-model <true|false>] [--split-max-tensors <NUM>] [--split-max-size <SIZE>]"
+            echo "    $0 --model <MODEL_ID> --quant-method <QUANT_METHOD> [--use-imatrix] [--output-filename <FILENAME>] [--split-model] [--split-max-tensors <NUM>] [--split-max-size <SIZE>] [--token-embedding-type <QUANT_METHOD>] [--leave-output-tensor] [--output-tensor-type <QUANT_METHOD>]"
             echo ""
             echo "  Short options:"
-            echo "    $0 -m <MODEL_ID> -q <QUANT_METHOD> [-i <true|false>] [-o <FILENAME>] [-s <true|false>] [-t <NUM>] [-z <SIZE>]"
+            echo "    $0 -m <MODEL_ID> -q <QUANT_METHOD> [-imatrix] [-o <FILENAME>] [-split]"
             echo ""
             echo "Examples:"
             echo "  $0 --model meta-llama/Llama-2-7b --quant-method Q4_K_M"
-            echo "  $0 -m meta-llama/Llama-2-7b -q Q4_K_M -i true"
-            echo "  $0 --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix true --output-filename Llama-2-7b-Q4_K_M.gguf"
-            echo "  $0 -m meta-llama/Llama-2-7b -q Q4_K_M -i true -o Llama-2-7b-Q4_K_M.gguf -s true -t 256 -z 4G"
+            echo "  $0 -m meta-llama/Llama-2-7b -q Q4_K_M -imatrix"
+            echo "  $0 --model meta-llama/Llama-2-7b --quant-method Q4_K_M --use-imatrix --output-filename Llama-2-7b-Q4_K_M.gguf"
+            echo "  $0 -m meta-llama/Llama-2-7b -q Q4_K_M -imatrix -o Llama-2-7b-Q4_K_M.gguf -split --split-max-tensors 256 --split-max-size 4G"
             exit 0
             ;;
         *)
@@ -237,26 +258,36 @@ QUANTIZE_CMD=(
     llama-quantize
 )
 
-# Add optional quantization flags
-# Note: The original script has logic for --leave-output-tensor vs --output-tensor-type
-# and --token-embedding-type. This script omits these for simplicity.
-# You can add them back if needed, but they require more input arguments.
-
 if [ "$USE_IMATRIX" = "true" ] && [ -f "$IMATRIX_FILE_PATH" ]; then
     QUANTIZE_CMD+=(
         --imatrix "$IMATRIX_FILE_PATH"
-        "$FP16_MODEL_PATH"
-        "$QUANTIZED_MODEL_PATH"
-        "$QUANT_METHOD"
     )
-else
+fi
+
+if [ -n "$TOKEN_EMBEDDING_TYPE" ]; then
     QUANTIZE_CMD+=(
-        "$FP16_MODEL_PATH"
-        "$QUANTIZED_MODEL_PATH"
-        "$QUANT_METHOD"
+        --token-embedding-type "$TOKEN_EMBEDDING_TYPE"
     )
 fi
 
+if [ "$LEAVE_OUTPUT_TENSOR" = "true" ]; then
+    QUANTIZE_CMD+=(
+        --leave-output-tensor
+    )
+else
+    if [ -n "$OUTPUT_TENSOR_TYPE" ]; then
+        QUANTIZE_CMD+=(
+            --output-tensor-type "$OUTPUT_TENSOR_TYPE"
+        )
+    fi
+fi
+
+QUANTIZE_CMD+=(
+    "$FP16_MODEL_PATH"
+    "$QUANTIZED_MODEL_PATH"
+    "$QUANT_METHOD"
+)
+
 echo "Running command: ${QUANTIZE_CMD[*]}"
 "${QUANTIZE_CMD[@]}"