Upload app.py

app.py

@@ -5,6 +5,8 @@ from transformers import AutoTokenizer
 from model import TransformerModel
 import gradio as gr
 
+from torch.ao.quantization.qconfig import float_qparams_weight_only_qconfig
+
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
 
@@ -22,32 +24,20 @@ def load_quantized_model(checkpoint_path):
         tie_word_embeddings=True,
     )
     
-    # Dynamic quant for embeddings
-    model.embed_tokens = torch.quantization.quantize_dynamic(
-        model.embed_tokens, {nn.Embedding}, dtype=torch.qint8
-    )
-    model.embed_positions = torch.quantization.quantize_dynamic(
-        model.embed_positions, {nn.Embedding}, dtype=torch.qint8
-    )
-    
-    # Static quant config for the rest of the model
-    model.qconfig = torch.quantization.get_default_qconfig("fbgemm")  # CPU
-    model = torch.quantization.prepare(model, inplace=False)
+    # This qconfig is typically for your other layers
+    default_qconfig = torch.quantization.get_default_qconfig("fbgemm")
+    model.qconfig = default_qconfig
+
+    # For embeddings, force the specialized config:
+    model.embed_tokens.qconfig = float_qparams_weight_only_qconfig
+    model.embed_positions.qconfig = float_qparams_weight_only_qconfig
 
-    #
-    # >>> RUN CALIBRATION HERE (forward pass with sample data) <<<
-    # e.g. with torch.no_grad():
-    #         for input_ids in some_calibration_loader:
-    #             outputs = model(input_ids)
-    #
+    # Then prepare, calibrate, and convert
+    model = torch.quantization.prepare(model, inplace=False)
 
+    # Calibration pass here...
     model = torch.quantization.convert(model, inplace=False)
-    
-    # Load checkpoint
-    checkpoint = torch.load(checkpoint_path, map_location="cpu")
-    model.load_state_dict(checkpoint)
 
-    model.eval()
     return model