app.py

import gradio as gr
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer, AutoModel
import tempfile
from huggingface_hub import HfApi
from huggingface_hub import list_models
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from packaging import version
import os
import spaces

MAP_QUANT_TYPE_TO_NAME = {
    "int4_weight_only": "int4wo", "int8_weight_only": "int8wo", "int8_dynamic_activation_int8_weight": "int8da8w"
}

def hello(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
    # ^ expect a gr.OAuthProfile object as input to get the user's profile
    # if the user is not logged in, profile will be None
    if profile is None:
        return "Hello !"
    return f"Hello {profile.name} !"

def check_model_exists(oauth_token: gr.OAuthToken | None, username, quantization_type, group_size, model_name, quantized_model_name):
    """Check if a model exists in the user's Hugging Face repository."""
    try:
        models = list_models(author=username, token=oauth_token.token)
        model_names = [model.id for model in models]
        if quantized_model_name : 
            repo_name = f"{username}/{quantized_model_name}"
        else : 
            if quantization_type == "int4_weight_only" : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}-gs{group_size}"
            else : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}"

        if repo_name in model_names:
            return f"Model '{repo_name}' already exists in your repository."
        else:
            return None  # Model does not exist
    except Exception as e:
        return f"Error checking model existence: {str(e)}"

def create_model_card(model_name, quantization_type, group_size):
    model_card = f"""---
base_model:
- {model_name}
---

# {model_name} (Quantized)

## Description
This model is a quantized version of the original model `{model_name}`. It has been quantized using {quantization_type} quantization with torchao.

## Quantization Details
- **Quantization Type**: {quantization_type}
- **Group Size**: {group_size if quantization_type == "int4_weight_only" else None}

## Usage
You can use this model in your applications by loading it directly from the Hugging Face Hub:

```python
from transformers import AutoModel

model = AutoModel.from_pretrained("{model_name}")"""
    
    return model_card

def load_model(model_name, quantization_config, auth_token) : 
    return AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=quantization_config, device_map="cpu", use_auth_token=auth_token.token)

def load_model_cpu(model_name, quantization_config, auth_token) : 
    return AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)

def quantize_model(model_name, quantization_type, group_size=128, auth_token=None, username=None):
    print(f"Quantizing model: {quantization_type}")
    if quantization_type == "int4_weight_only" : 
        quantization_config = TorchAoConfig(quantization_type, group_size=group_size)
    else : 
        quantization_config = TorchAoConfig(quantization_type)
    model = load_model(model_name, quantization_config=quantization_config, auth_token=auth_token)

    return model

def save_model(model, model_name, quantization_type, group_size=128, username=None, auth_token=None, quantized_model_name=None):
    print("Saving quantized model")
    with tempfile.TemporaryDirectory() as tmpdirname:


        model.save_pretrained(tmpdirname, safe_serialization=False, use_auth_token=auth_token.token)
        if quantized_model_name : 
            repo_name = f"{username}/{quantized_model_name}"
        else : 
            if quantization_type == "int4_weight_only" : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}-gs{group_size}"
            else : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}"

        model_card = create_model_card(repo_name, quantization_type, group_size)
        with open(os.path.join(tmpdirname, "README.md"), "w") as f:
            f.write(model_card)
        # Push to Hub
        api = HfApi(token=auth_token.token)
        api.create_repo(repo_name, exist_ok=True)
        api.upload_folder(
            folder_path=tmpdirname,
            repo_id=repo_name,
            repo_type="model",
        )
    return f'<h1> 🤗 DONE</h1><br/>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a>'

def quantize_and_save(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None, model_name, quantization_type, group_size, quantized_model_name):
    if oauth_token is None : 
        return "Error : Please Sign In to your HuggingFace account to use the quantizer"
    if not profile:
        return "Error: Please Sign In to your HuggingFace account to use the quantizer"
    exists_message = check_model_exists(oauth_token, profile.username, quantization_type, group_size, model_name, quantized_model_name)
    if exists_message : 
        return exists_message
    if quantization_type == "int4_weight_only" : 
        return "int4_weight_only not supported on cpu"
    if not group_size.isdigit() :
        return "group_size must be a number"
    
    group_size = int(group_size)

    try:
        quantized_model = quantize_model(model_name, quantization_type, group_size, oauth_token, profile.username)
        return save_model(quantized_model, model_name, quantization_type, group_size, profile.username, oauth_token, quantized_model_name)
    except Exception as e : 
        return e


css="""/* Custom CSS to allow scrolling */
.gradio-container {overflow-y: auto;}
"""
with gr.Blocks(theme=gr.themes.Ocean(), css=css) as app:
    gr.Markdown(
        """
        # 🤗 LLM Model TorchAO Quantization App
        
        Quantize your favorite Hugging Face models using TorchAO and save them to your profile!
        """
    )

    gr.LoginButton(elem_id="login-button", elem_classes="center-button", min_width=250)

    m1 = gr.Markdown()
    app.load(hello, inputs=None, outputs=m1)


    radio = gr.Radio(["show", "hide"], label="Show Instructions", value="hide")
    instructions = gr.Markdown(
        """
        ## Instructions
        1. Login to your HuggingFace account
        2. Enter the name of the Hugging Face LLM model you want to quantize (Make sure you have access to it)
        3. Choose the quantization type.
        4. Optionally, specify the group size.
        5. Optionally, choose a custom name for the quantized model
        6. Click "Quantize and Save Model" to start the process.
        7. Once complete, you'll receive a link to the quantized model on Hugging Face.
        
        Note: This process may take some time depending on the model size and your hardware you can check the container logs to see where are you at in the process!
        """,
        visible=False
    )
    def update_visibility(radio):  
        value = radio 
        if value == "show":
            return gr.Textbox(visible=True) 
        else:
            return gr.Textbox(visible=False)
    radio.change(update_visibility, radio, instructions)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                model_name = HuggingfaceHubSearch(
                    label="Hub Model ID",
                    placeholder="Search for model id on Huggingface",
                    search_type="model",
                    scale=2
                )
            with gr.Row():
                with gr.Column():
                    quantization_type = gr.Dropdown(
                        info="Quantization Type",
                        choices=["int4_weight_only", "int8_weight_only", "int8_dynamic_activation_int8_weight"],
                        value="int8_weight_only",
                        filterable=False,
                        show_label=False,
                    )
                    group_size = gr.Textbox(
                        info="Group Size (only for int4_weight_only)",
                        value=128,
                        interactive=True,
                        show_label=False
                    )
                    quantized_model_name = gr.Textbox(
                        info="Model Name (optional : to override default)",
                        value="",
                        interactive=True,
                        show_label=False
                    )
        with gr.Column():
            quantize_button = gr.Button("Quantize and Save Model", variant="primary")
            output_link = gr.Markdown(label="Quantized Model Link", container=True, min_height=40)
    
    
    # Adding CSS styles for the username box
    app.css = """
    #username-box {
        background-color: #f0f8ff; /* Light color */
        border-radius: 8px;
        padding: 10px;
    }
    """
    app.css = """
    .center-button {
        display: flex;
        justify-content: center;
        align-items: center;
        margin: 0 auto; /* Center horizontally */
    }
    """
    
    quantize_button.click(
        fn=quantize_and_save,
        inputs=[model_name, quantization_type, group_size, quantized_model_name],
        outputs=[output_link]
    )


# Launch the app
app.launch()