Update app.py

app.py

@@ -119,13 +119,14 @@ def clean_up(cache, origin_len):
         cache.value_cache[i] = cache.value_cache[i][:, :, :origin_len, :]
     return cache
 
-def get_cache_memory(self):
+def calculate_cache_size(cache):
+    """Calculate the total memory used by the key-value cache in bytes."""
     total_memory = 0
-    for key in self.key_cache:
+    for key in cache.key_cache:
         total_memory += key.element_size() * key.nelement()
-    for value in self.value_cache:
+    for value in cache.value_cache:
         total_memory += value.element_size() * value.nelement()
-    return total_memory/(1024*1024)
+    return total_memory /(1024*1024)
     
 @st.cache_resource
 def load_model_and_tokenizer(doc_text_count):
@@ -198,7 +199,7 @@ if uploaded_file:
         st.text(doc_text[:500] + "..." if len(doc_text) > 500 else doc_text)
     query = st.text_input("🔎 Ask a question about the document:")
     if query and st.button("Generate Answer"):
-        with st.spinner("Generating answer... ()"):
+        with st.spinner("Generating answer..."):
             current_cache = clone_cache(cache)
             t_clone_end = time()
             Cache_create_time = t_clone_end - t1
@@ -213,7 +214,7 @@ if uploaded_file:
             response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
             t_gen_end = time()
             last_generation_time = t_gen_end - t_gen_start
-            cache_mem_bytes = cache.get_cache_memory()
+            cache_mem_bytes = calculate_cache_size(cache)
             st.success("Answer:")
             st.write(response)
             st.info(f"Cache create Time: {Cache_create_time:.2f} s  |  Generation Time: {last_generation_time:.2f} s   ")