Update app.py

app.py

@@ -4,6 +4,7 @@
 import os
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, pipeline
 import streamlit as st
+import torch
 
 # Define the model repository
 # REPO_NAME = 'schuler/experimental-JP47D20'
@@ -12,7 +13,6 @@ REPO_NAME = 'schuler/experimental-JP47D21-KPhi-3-micro-4k-instruct'
 # Configure the Streamlit app
 st.set_page_config(page_title="Experimental KPhi3 Model - Currently in Training", page_icon="🤗")
 st.title("Experimental KPhi3 Model - Currently in Training")
-st.markdown(f"*This chat uses the {REPO_NAME} model. Feel free to ask questions such as 'What is biology?' or 'What is the human body?'*")
 
 # Load tokenizer and model
 @st.cache_resource(show_spinner="Loading model...")
@@ -24,6 +24,9 @@ def load_model(repo_name):
 
 tokenizer, generator_conf, model = load_model(REPO_NAME)
 
+st.markdown(f"*This chat uses the {REPO_NAME} model with {model.get_memory_footprint() / 1e6:.2f} MB memory footprint.")
+st.markdown(f"You may ask questions such as 'What is biology?' or 'What is the human body?'*")
+
 try:
     generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 except Exception as e: