Update app.py

app.py

@@ -1,4 +1,30 @@
 import gradio as gr
+from huggingface_hub import snapshot_download
+from accelerate.utils import BnbQuantizationConfig
+from accelerate.utils import load_and_quantize_model
+from accelerate import Accelerator
 
-def quantize():
-    
+model_path="marcsun13/gpt2-xl-linear-sharded"
+
+def quantize(model_path=model_path):
+    print("1")
+    weights_location = snapshot_download(repo_id=f"{model_path}")
+    print("2")
+    bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)
+    #bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
+    print("3")
+    quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
+    print("4")
+    accelerate = Accelerator()
+    print("5")
+    new_weights_location = "./model"
+    print("6")
+    accelerate.save_model(quantized_model, new_weights_location)
+    print("7")
+    quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
+    print("Done")
+
+with gr.Blocks() as app:
+    btn=gr.Button()
+    btn.click(quantize,None,None)
+app.launch()