import gradio as gr import torch from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer import tempfile from huggingface_hub import HfApi from huggingface_hub import list_models from packaging import version import os 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-{quantization_type.lower()}-gs_{group_size}" else : repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{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 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 = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.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_card = create_model_card(model_name, quantization_type, group_size) with open(os.path.join(tmpdirname, "README.md"), "w") as f: f.write(model_card) 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-{quantization_type.lower()}-gs_{group_size}" else : repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{quantization_type.lower()}" # Push to Hub api = HfApi() api.create_repo(repo_name, exist_ok=True) api.upload_folder( folder_path=tmpdirname, repo_id=repo_name, repo_type="model", ) return f"https://huggingface.co./{repo_name}" 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 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) with gr.Blocks(theme=gr.themes.Soft()) as app: gr.Markdown( """ # 🚀 Model Quantization App Quantize your favorite Hugging Face models and save them to your profile! """ ) gr.LoginButton(elem_id="login-button", elem_classes="center-button") m1 = gr.Markdown() app.load(hello, inputs=None, outputs=m1) with gr.Row(): with gr.Column(): model_name = gr.Textbox( label="Model Name", placeholder="e.g., meta-llama/Meta-Llama-3-8B", value="meta-llama/Meta-Llama-3-8B" ) quantization_type = gr.Dropdown( label="Quantization Type", choices=["int4_weight_only", "int8_weight_only", "int8_dynamic_activation_int8_weight"], value="int8_weight_only" ) group_size = gr.Number( label="Group Size (only for int4_weight_only)", value=128, interactive=True ) quantized_model_name = gr.Textbox( label="Model Name (optional : to override default)", value="", interactive=True ) # with gr.Row(): # username = gr.Textbox( # label="Hugging Face Username", # placeholder="Enter your Hugging Face username", # value="", # interactive=True, # elem_id="username-box" # ) with gr.Column(): quantize_button = gr.Button("Quantize and Save Model", variant="primary") output_link = gr.Textbox(label="Quantized Model Link") gr.Markdown( """ ## Instructions 1. Enter the name of the Hugging Face model you want to quantize. 2. Choose the quantization type. 3. Optionally, specify the group size. 4. Click "Quantize and Save Model" to start the process. 5. 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. """ ) # 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()