Update app.py

app.py

@@ -29,7 +29,7 @@ from utils import (
 )
 
 # Initialize the model and tokenizer.
-api_token = os.getenv("HF_TOKEN")
+api_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
 model_name = "meta-llama/Llama-3.1-8B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_token)
 model = AutoModelForCausalLM.from_pretrained(model_name, token=api_token, torch_dtype=torch.float16)
@@ -591,7 +591,7 @@ def update_token_breakdown(token_count, retrieval_slider, global_local_value):
 
     rag_tokens = int(retrieval_context_length * (1.0 - (percentage / 100)))
     kv_tokens = retrieval_context_length - rag_tokens
-    return f"Token Breakdown: {kv_tokens} tokens (KV compression), {rag_tokens} tokens (RAG retrieval)"
+    return f"Token Breakdown: {kv_tokens} tokens (KV compression), {rag_tokens} tokens (RAG retrieval)", f"Number of tokens after compression: {retrieval_context_length}"
 
 ##########################################################################
 # Gradio Interface
@@ -797,12 +797,12 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(font=["Arial", gr.themes.GoogleFont
             retrieval_slider.change(
                 fn=update_token_breakdown,
                 inputs=[hidden_token_count, retrieval_slider, global_local_slider],
-                outputs=tokens_breakdown_text
+                outputs=[tokens_breakdown_text, retrieval_info_text]
             )
             global_local_slider.change(
                 fn=update_token_breakdown,
                 inputs=[hidden_token_count, retrieval_slider, global_local_slider],
-                outputs=tokens_breakdown_text
+                outputs=[tokens_breakdown_text, retrieval_info_text]
             )
             
             # Compress button: Prepare compression and then update chat status.