Update SSE_quantize.py

SSE_quantize.py

@@ -187,16 +187,19 @@ class SSEQ(InputModule):
         with open(bin_path, "rb") as f:
             raw = f.read()
 
-        vocab = state["dyt.alpha"].shape[0]  # hidden dim
         hidden = state["dyt.alpha"].shape[0]
+        total_uint8 = len(raw)
+
+        bytes_per_row = hidden // 2 + 4
+        vocab = total_uint8 // bytes_per_row
 
         packed_size = vocab * hidden // 2
 
         packed = np.frombuffer(raw[:packed_size], dtype=np.uint8)
         scales = np.frombuffer(raw[packed_size:], dtype=np.float32)
 
-        packed = packed.reshape(vocab, hidden//2)
-        scales = scales.reshape(-1, 1)
+        packed = packed.reshape(vocab, hidden // 2)
+        scales = scales.reshape(vocab, 1)
 
         emb = dequantize_q4_k_m(packed, scales)