Upload app.py

app.py

@@ -23,19 +23,21 @@ def load_quantized_model(checkpoint_path):
         tie_word_embeddings=True,
     )
     
-    # Set quantization config for ALL embedding layers
-    for name, module in model.named_modules():
-        if isinstance(module, nn.Embedding):  # Now works because `nn` is imported
-            module.qconfig = quantization.float_qparams_weight_only_qconfig
+    # Dynamic quantization for embeddings
+    model.embed_tokens = torch.ao.quantization.quantize_dynamic(
+        model.embed_tokens, {nn.Embedding}, dtype=torch.qint8
+    )
+    model.embed_positions = torch.ao.quantization.quantize_dynamic(
+        model.embed_positions, {nn.Embedding}, dtype=torch.qint8
+    )
     
-    # Apply static quantization to the rest of the model
-    model.qconfig = quantization.default_qconfig
-    model = quantization.prepare(model, inplace=False)
-    model = quantization.convert(model, inplace=False)
+    # Static quantization for other layers
+    model.qconfig = torch.ao.quantization.default_qconfig
+    model = torch.ao.quantization.prepare(model, inplace=False)
+    model = torch.ao.quantization.convert(model, inplace=False)
     
-    # Load the checkpoint
+    # Load checkpoint
     checkpoint = torch.load(checkpoint_path, map_location="cpu")
-    # model.load_state_dict(checkpoint["model_state_dict"])
     model.load_state_dict(checkpoint)
     
     model.eval()