Update app.py

app.py

@@ -1,77 +1,9 @@
-import streamlit as st
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers.cache_utils import DynamicCache
-import os
-from time import time
 import pandas as pd
-import psutil
-
-# ==============================
-# Helper: Human-readable bytes
-def sizeof_fmt(num, suffix="B"):
-    # Formats bytes as human-readable (e.g. 1.5 GB)
-    for unit in ["", "K", "M", "G", "T"]:
-        if abs(num) < 1024.0:
-            return f"{num:3.2f} {unit}{suffix}"
-        num /= 1024.0
-    return f"{num:.2f} P{suffix}"
-
-# ==============================
-# System & Cache Resource Stats
-def get_system_stats(doc_text=None, cache_mem_bytes=0):
-    ram = psutil.virtual_memory()
-    cpu = psutil.cpu_percent()
-    disk = psutil.disk_usage('/')
-    used, total = ram.used, ram.total
-    stats = {
-        "Input Tokens": st.session_state.input_tokens_count,
-        "Output Tokens": st.session_state.output_tokens_count,
-        "Generated Tokens": st.session_state.generated_tokens_count,
-        "Document Size (chars)": len(doc_text) if doc_text else 0,
-        "Document Size (KB)": f"{len(doc_text.encode('utf-8')) / 1024:.2f}" if doc_text else 0,
-    }
-    if torch.cuda.is_available():
-        gpu_mem_alloc = torch.cuda.memory_allocated()
-        gpu_mem_total = torch.cuda.get_device_properties(0).total_memory
-        stats["GPU Used"] = sizeof_fmt(gpu_mem_alloc)
-        stats["GPU Total"] = sizeof_fmt(gpu_mem_total)
-        stats["GPU Usage (%)"] = round(100 * gpu_mem_alloc / gpu_mem_total, 2) if gpu_mem_total else 0
-    else:
-        stats["GPU Used"] = "N/A"
-        stats["GPU Total"] = "N/A"
-        stats["GPU Usage (%)"] = "N/A"
-
-    stats["KV Cache Memory Used"] = sizeof_fmt(cache_mem_bytes)
-    stats["KV Cache as % RAM"] = f"{(cache_mem_bytes / total) * 100:.2f}%" if total > 0 else "N/A"
-    stats["KV Cache as % GPU"] = (
-        f"{(cache_mem_bytes / torch.cuda.get_device_properties(0).total_memory) * 100:.2f}%"
-        if torch.cuda.is_available() else "N/A"
-    )
-    return stats
-
-def cache_stats_table(cache):
-    if cache is None:
-        return pd.DataFrame(), 0
-    rows = []
-    total_mem = 0
-    for i, (key, value) in enumerate(zip(cache.key_cache, cache.value_cache)):
-        key_mem = key.element_size() * key.nelement()
-        value_mem = value.element_size() * value.nelement()
-        total_mem += key_mem + value_mem
-        row = {
-            "Layer": i,
-            "Key Shape": str(tuple(key.shape)),
-            "Value Shape": str(tuple(value.shape)),
-            "Total Mem": sizeof_fmt(key_mem + value_mem),
-            "Last Key Tokens": str(tuple(key[..., -1:, :].shape)),
-            "Last Value Tokens": str(tuple(value[..., -1:, :].shape)),
-        }
-        rows.append(row)
-    return pd.DataFrame(rows), total_mem
+import streamlit as st
+from time import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ==============================
-# Core Model and Caching Logic
 def generate(model, input_ids, past_key_values, max_new_tokens=50):
     """Token-by-token generation using cache for speed."""
     device = model.model.embed_tokens.weight.device
@@ -163,25 +95,6 @@ def load_document_and_cache(file_path):
         st.error(f"Document file not found at {file_path}")
         return None, None, None, None
 
-# ==============================
-# Main Streamlit UI and Workflow
-st.title("DeepSeek QA: Supercharged Caching & Memory Dashboard")
-
-# Initialize session state variables
-if "doc_uploaded" not in st.session_state:
-    st.session_state.update({
-        "doc_uploaded": False,
-        "doc_text": None,
-        "cache": None,
-        "origin_len": None,
-        "doc_text_count": None,
-        "generated_tokens_count": 0,
-        "input_tokens_count": 0,
-        "output_tokens_count": 0,
-        "cache_gen_duration": 0.0,
-        "output_gen_duration": 0.0,
-    })
-
 # File upload
 uploaded_file = st.file_uploader("📝 Upload your document (.txt)", type="txt")
 if uploaded_file:
@@ -201,13 +114,39 @@ if uploaded_file:
         "output_tokens_count": 0,
         "cache_gen_duration": 0.0,
         "output_gen_duration": 0.0,
+        "stats_table": pd.DataFrame(columns=[
+            "Timestamp", 
+            "Document Characters", 
+            "Document Size (KB)", 
+            "Cache Memory Used", 
+            "Cache as % RAM", 
+            "Cache as % GPU"
+        ])
     })
 
     # Measure cache generation time
     cache_start_time = time()
-    st.session_state["cache"], st.session_state["origin_len"], st.session_state["doc_text"], st.session_state["doc_text_count"] = load_document_and_cache(temp_file_path)
-    cache_end_time = time()
-    st.session_state["cache_gen_duration"] = cache_end_time - cache_start_time
+    try:
+        # Simulated function `load_document_and_cache`
+        st.session_state["cache"], st.session_state["origin_len"], st.session_state["doc_text"], st.session_state["doc_text_count"] = load_document_and_cache(temp_file_path)
+        cache_end_time = time()
+        st.session_state["cache_gen_duration"] = cache_end_time - cache_start_time
+
+        # Log the resource usage after cache generation
+        log_time = time()
+        total_ram = torch.cuda.get_device_properties(0).total_memory if torch.cuda.is_available() else 1
+        cache_mem_bytes = st.session_state["cache"].get_cache_memory() if st.session_state["cache"] else 0
+        st.session_state["stats_table"] = log_resource_usage(
+            st.session_state["stats_table"],
+            st.session_state["doc_text"],
+            cache_mem_bytes,
+            total_ram,
+            log_time
+        )
+
+    except Exception as e:
+        st.error(f"Failed to generate cache: {str(e)}")
+        st.stop()
 
     # Display document preview
     with st.expander("📄 Document Preview"):
@@ -218,6 +157,7 @@ if uploaded_file:
     if query and st.button("Generate Answer"):
         with st.spinner("Generating answer..."):
             try:
+                # Simulated function `load_model_and_tokenizer`
                 model, tokenizer = load_model_and_tokenizer(st.session_state["doc_text_count"])
                 current_cache = clone_cache(st.session_state["cache"])
 
@@ -226,7 +166,10 @@ if uploaded_file:
                 full_prompt = f"<|user|>\nQuestion: {query}\n<|assistant|>"
                 input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids
                 st.session_state["input_tokens_count"] += input_ids.shape[-1]
+
+                # Simulated function `generate`
                 output_ids, generated_tokens_count = generate(model, input_ids, current_cache)
+
                 response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
                 output_gen_end_time = time()
                 st.session_state["output_gen_duration"] = output_gen_end_time - output_gen_start_time
@@ -242,15 +185,22 @@ if uploaded_file:
                     f"Output Generation Time: {st.session_state['output_gen_duration']:.2f} seconds"
                 )
 
-                # Display cache stats table
-                table, total_mem = cache_stats_table(st.session_state["cache"])
-                if not table.empty:
-                    st.write("Cache Statistics Table:")
-                    st.dataframe(table)
-                else:
-                    st.write("No data available in the cache.")
+                # Log the resource usage after output generation
+                log_time = time()
+                cache_mem_bytes = st.session_state["cache"].get_cache_memory() if st.session_state["cache"] else 0
+                st.session_state["stats_table"] = log_resource_usage(
+                    st.session_state["stats_table"],
+                    st.session_state["doc_text"],
+                    cache_mem_bytes,
+                    total_ram,
+                    log_time
+                )
 
             except Exception as e:
                 st.error(f"Failed to generate answer: {str(e)}")
+
+    # Display the stats table
+    st.write("📊 Resource Usage Log:")
+    st.dataframe(st.session_state["stats_table"])
 else:
     st.info("Please upload a document to start.")