Upload app.py

app.py

@@ -1,6 +1,6 @@
 import torch
-import torch.nn as nn  # Added missing import
-import torch.ao.quantization as quantization
+import torch.nn as nn
+import torch.quantization  # <--- Use the older namespace for default_qconfig
 from transformers import AutoTokenizer
 from model import TransformerModel
 import gradio as gr
@@ -9,7 +9,6 @@ import gradio as gr
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
 
 def load_quantized_model(checkpoint_path):
-    # Define the model architecture
     model = TransformerModel(
         vocab_size=49152,
         hidden_size=576,
@@ -23,29 +22,37 @@ def load_quantized_model(checkpoint_path):
         tie_word_embeddings=True,
     )
     
-    # Dynamic quantization for embeddings
-    model.embed_tokens = torch.ao.quantization.quantize_dynamic(
+    # Dynamic quant for embeddings
+    model.embed_tokens = torch.quantization.quantize_dynamic(
         model.embed_tokens, {nn.Embedding}, dtype=torch.qint8
     )
-    model.embed_positions = torch.ao.quantization.quantize_dynamic(
+    model.embed_positions = torch.quantization.quantize_dynamic(
         model.embed_positions, {nn.Embedding}, dtype=torch.qint8
     )
     
-    # 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)
+    # 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)
+
+    #
+    # >>> 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)
+    #
+
+    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
 
 
 # Load the quantized model
-model = load_quantized_model("checkpoint_quantized.pt")
+model = load_quantized_model("quantized_model.pt")
 
 # Function to generate text
 def generate_text(prompt, max_length=50, temperature=1.0, top_k=50):