update code

app.py

@@ -8,7 +8,7 @@ from PyPDF2 import PdfReader
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
-from time import sleep
+import time
 
 # notes
 # https://huggingface.co/docs/transformers/pad_truncation
@@ -118,7 +118,7 @@ def main():
                     legacy=False,
                     model_max_length=1000,
                     trust_remote_code=True, 
-                    cache_dir="model_cache"
+                    #cache_dir="model_cache"
                 )
                 base_model = "model_cache/models--ccdv--lsg-bart-base-16384-pubmed/snapshots/4072bc1a7a94e2b4fd860a5fdf1b71d0487dcf15"
                 #base_model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -134,7 +134,7 @@ def main():
                     truncation=True,
                     legacy=False,
                     model_max_length=1000, 
-                    cache_dir="model_cache"
+                    #cache_dir="model_cache"
                 )
                 base_model = "model_cache/models--MBZUAI--LaMini-Flan-T5-77M/snapshots/c5b12d50a2616b9670a57189be20055d1357b474"
                 #base_model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -169,12 +169,17 @@ def main():
                 )
                 pdf_viewer = displayPDF(filepath)
             with col2:
+                start = time.time()
                 with st.spinner("Summarizing..."):
                     summary = llm_pipeline(tokenizer, base_model, input_text)
                     postproc_text_length = postproc_count(summary)
+                end = time.time()
+                duration = end - start
                 st.info(
                     "PDF Summary  |  Number of words: "
                     f"{postproc_text_length:,}"
+                    + "  |  Summarization time: "
+                    f"{duration:.0f}" + " seconds"
                 )
                 st.success(summary)