Update app.py

An attempt at avx512_vnni int8 quanitzation using ONNX runtime.

app.py

@@ -1,30 +1,66 @@
 
+# import gc
+
+# import gradio as gr
+# import torch
+# from transformers import AutoTokenizer, AutoModelForCausalLM, HqqConfig
+
+# quant_config = HqqConfig(nbits=8, group_size=64)
+
+# MODEL_ID = "HuggingFaceTB/SmolLM3-3B"
+# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# print("Loading tokenizer & model…")
+# tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# # model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to(DEVICE)
+
+# model =\
+#         AutoModelForCausalLM\
+#                 .from_pretrained(
+#                         MODEL_ID, 
+#                         torch_dtype=torch.float16, 
+#                         # device_map="cuda", 
+#                         quantization_config=quant_config
+#         ).to(DEVICE)
+
+#gc.collect()
+
+#########
+
 import gc
 
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, HqqConfig
+from transformers import AutoTokenizer
+from optimum.onnxruntime import ORTModelForCausalLM, ORTQuantizer
+from optimum.onnxruntime.configuration import AutoQuantizationConfig
+
+MODEL_NAME = "HuggingFaceTB/SmolLM3-3B"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-quant_config = HqqConfig(nbits=8, group_size=64)
+qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=True)
+quantizer = ORTQuantizer.from_pretrained(MODEL_NAME)
 
-MODEL_ID = "HuggingFaceTB/SmolLM3-3B"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Step 4: Perform quantization saving output in a new directory
+quantized_model_dir = "./quantized_model"
+print("Starting quantization...")
+quantizer.quantize(save_dir=quantized_model_dir, quantization_config=qconfig)
 
-print("Loading tokenizer & model…")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-# model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to(DEVICE)
 
-model =\
-        AutoModelForCausalLM\
-                .from_pretrained(
-                        MODEL_ID, 
-                        torch_dtype=torch.float16, 
-                        # device_map="cuda", 
-                        quantization_config=quant_config
-        ).to(DEVICE)
+del(quantizer)
+del(qconfig)
 
+# Run garbage collection again to release memory from quantizer objects
 gc.collect()
 
+# Step 5: Load the quantized ONNX model for inference
+print("Loading quantized ONNX model for inference...")
+model = ORTModelForCausalLM.from_pretrained(quantized_model_dir)
+
+# Garbage collection again after final loading
+gc.collect()
+
+
 #########
 
 # print("Loading tokenizer & model…")