Update app.py

app.py

@@ -6,7 +6,20 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from functools import partial
 import gc
 
+def get_model_size(model: nn.Module, data_width=16, group_size=-1):
 
+    if group_size != -1:
+        data_width += (16 + 4) / group_size
+
+    num_elements = 0
+    for param in model.parameters():
+        num_elements += param.numel()
+    return num_elements * data_width
+
+Byte = 8
+KiB = 1024 * Byte
+MiB = 1024 * KiB
+GiB = 1024 * MiB
 
 # core quantization method (simulated quantization)
 def pseudo_quantize_tensor(w, n_bit=4, q_group_size=-1):
@@ -62,7 +75,7 @@ model_path = "facebook/opt-125m"
 
 model_q_path = "facebook/opt-125m_3bit"
 
-
+model = AutoModelForCausalLM.from_pretrained(model_q_path, device_map="auto", offload_folder=offload_folder)
 generator = pipeline('text-generation', model="facebook/opt-125m_3bit")
 #generator_q = pipeline('text-generation', model="facebook/opt-125m-awq")
 
@@ -98,7 +111,8 @@ def generate_text_from_quantized(prompt):
     return generated_text
 '''
 # Create a Gradio interface
-iface = gr.Interface(fn=generate_text_pip, inputs="text", outputs="text")
+model_size = get_model_size(model, data_width=32, group_size=128)
+iface = gr.Interface(fn=generate_text_pip, inputs="text", outputs="text", description=f"model size: {model_size/MiB:.2f} MiB")
 
 #iface_2 = gr.Interface(fn=generate_text_pip_q, inputs="text", outputs="text")