# ============================================================
# Quark AWQ INT4 Quantization Script (PowerShell)
# Auto-generated by Claude Code /quantize skill
# ============================================================

# Activate conda environment
Write-Host "[INFO] Activating quark_0.11 conda environment..."
& "C:\ProgramData\miniforge3\Scripts\conda.exe" activate quark_0.11

# Set environment variables
$env:TORCHINDUCTOR_DISABLE = "1"
$env:TORCH_COMPILE_DISABLE = "1"

# Use Python from quark_0.11 environment
$PYTHON_EXE = "C:\ProgramData\miniforge3\envs\quark_0.11\python.exe"

# Check and store package versions for README
Write-Host "[INFO] Checking package versions..."
$PYTHON_VERSION = & $PYTHON_EXE -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')" 2>&1
$TORCH_VERSION = & $PYTHON_EXE -c "import torch; print(torch.__version__)" 2>&1
$TRANSFORMERS_VERSION = & $PYTHON_EXE -c "import transformers; print(transformers.__version__)" 2>&1
$QUARK_VERSION = & $PYTHON_EXE -c "import quark; print(quark.__version__)" 2>&1

Write-Host "[INFO] Python: $PYTHON_VERSION, Torch: $TORCH_VERSION, Transformers: $TRANSFORMERS_VERSION, Quark: $QUARK_VERSION"
Write-Host ""

# Configuration variables
$PYTHON_SCRIPT = "quantize_quark.py"
$MODEL_DIR = "microsoft/Phi-4-mini-instruct"
$OUTPUT_DIR = "phi4_gptq_grp128_exlmhead"
$DATA_TYPE = "bfloat16"
# Options: auto | float16 | bfloat16 | float32

$GROUP_SIZE = 128
# Options: 128 | 32
# This sets the quantization scheme to uint4_wo_<GROUP_SIZE>

$EXCLUDE_LMHEAD = $true
# Set to $true to exclude lm_head from quantization (uses template default: exclude_layers_name=["lm_head"])
# Set to $false to quantize lm_head (passes --exclude_layers flag to override template with empty list)

# ============================================================
# Run quantization
# ============================================================

Write-Host ""
Write-Host "[INFO] Starting quantization..."
Write-Host "  Model      : $MODEL_DIR"
Write-Host "  Output dir : $OUTPUT_DIR"
Write-Host "  Data type  : $DATA_TYPE"
Write-Host "  Group size : $GROUP_SIZE"
Write-Host "  Exclude lm_head: $EXCLUDE_LMHEAD"
Write-Host ""

# Build quant_scheme from GROUP_SIZE
$QUANT_SCHEME = "uint4_wo_$GROUP_SIZE"

# Build command arguments
$args = @(
    "--model_dir", "$MODEL_DIR",
    "--output_dir", "$OUTPUT_DIR",
    "--device", "cpu",
    "--quant_scheme", "$QUANT_SCHEME",
    "--num_calib_data", "128",
    "--seq_len", "512",
    "--quant_algo", "gptq",
    "--dataset", "pileval_for_awq_benchmark",
    "--model_export", "hf_format",
    "--data_type", "$DATA_TYPE"
)

# Add --exclude_layers flag to override template (quantize lm_head)
if (-not $EXCLUDE_LMHEAD) {
    $args += "--exclude_layers"
}

# Run Python with output displayed on screen AND saved to log file
& $PYTHON_EXE $PYTHON_SCRIPT $args 2>&1 | Tee-Object -FilePath screen.dump.log

Write-Host ""
if ($LASTEXITCODE -eq 0) {
    Write-Host "[INFO] Quantization complete. Output saved to: $OUTPUT_DIR"
    Write-Host "[INFO] Full log saved to: screen.dump.log"

    # Copy files to output directory
    Write-Host ""
    Write-Host "[INFO] Copying files to output directory..."

    $LOG_FILE = "screen.dump.log"
    $PS_SCRIPT = $MyInvocation.MyCommand.Path

    # Convert screen.dump.log to UTF-8 before copying (for HuggingFace compatibility)
    Write-Host "[INFO] Converting $LOG_FILE to UTF-8..."
    $LOG_CONTENT = Get-Content $LOG_FILE -Encoding Unicode -Raw -ErrorAction SilentlyContinue
    if ($LOG_CONTENT) {
        $LOG_DEST = Join-Path $OUTPUT_DIR $LOG_FILE
        Set-Content -Path $LOG_DEST -Value $LOG_CONTENT -Encoding UTF8
        Write-Host "[INFO] Copied $LOG_FILE to $OUTPUT_DIR (converted to UTF-8)"
    } else {
        # Fallback: try regular copy if Unicode read fails
        Copy-Item $LOG_FILE -Destination $OUTPUT_DIR -Force
        Write-Host "[INFO] Copied $LOG_FILE to $OUTPUT_DIR"
    }

    Copy-Item $PYTHON_SCRIPT -Destination $OUTPUT_DIR -Force
    Copy-Item $PS_SCRIPT -Destination $OUTPUT_DIR -Force

    Write-Host "[INFO] Copied $PYTHON_SCRIPT to $OUTPUT_DIR"
    Write-Host "[INFO] Copied $(Split-Path $PS_SCRIPT -Leaf) to $OUTPUT_DIR"

    # Create README.md in output directory
    Write-Host ""
    Write-Host "[INFO] Creating README.md in output directory..."
    $README_PATH = Join-Path $OUTPUT_DIR "README.md"

    # Extract perplexity from screen.dump.log (last line)
    $LOG_PATH = Join-Path $OUTPUT_DIR $LOG_FILE
    $PERPLEXITY = "N/A"
    if (Test-Path $LOG_PATH) {
        $lastLine = Get-Content $LOG_PATH -Tail 1
        if ($lastLine -match "Perplexity:\s+([\d\.]+)") {
            $PERPLEXITY = $matches[1]
        }
    }

    # Build lm_head status text
    $LMHEAD_STATUS = if ($EXCLUDE_LMHEAD) { "excluded" } else { "included" }

    # Build Quark command (reconstruct from executed command)
    $QUARK_COMMAND = "python $PYTHON_SCRIPT --model_dir $MODEL_DIR --output_dir $OUTPUT_DIR --device cpu --quant_scheme $QUANT_SCHEME --num_calib_data 128 --seq_len 512 --quant_algo awq --dataset pileval_for_awq_benchmark --model_export hf_format --data_type $DATA_TYPE"
    if (-not $EXCLUDE_LMHEAD) {
        $QUARK_COMMAND += " --exclude_layers"
    }

    # Create README content
    $README_CONTENT = @"
# Quark Quantization Details

**Package Versions**:
- Python: $PYTHON_VERSION
- PyTorch: $TORCH_VERSION
- Transformers: $TRANSFORMERS_VERSION
- Quark: $QUARK_VERSION

**Quantization Setup**: AWQ, Group Size $GROUP_SIZE, $DATA_TYPE, lm_head $LMHEAD_STATUS

**Perplexity**: $PERPLEXITY

**Quark Command**:
``````
$QUARK_COMMAND
``````
"@

    Set-Content -Path $README_PATH -Value $README_CONTENT -Encoding UTF8
    Write-Host "[INFO] Created README.md in $OUTPUT_DIR (Perplexity: $PERPLEXITY)"
} else {
    Write-Host "[ERROR] Quantization failed with exit code $LASTEXITCODE"
    Write-Host "[INFO] Check screen.dump.log for details"
}

Write-Host ""
Write-Host "Press any key to continue..."
$null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")