Update app.py

app.py

@@ -1,12 +1,10 @@
 import streamlit as st
 import torch
-from transformers import  AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers.cache_utils import DynamicCache
 import os
 from time import time
 import pandas as pd
-
-import os
 from huggingface_hub import login
 
 HF_TOKEN = os.getenv("NEX_MODEL")  # Updated key name for clarity
@@ -14,7 +12,8 @@ HF_TOKEN = os.getenv("NEX_MODEL")  # Updated key name for clarity
 if not HF_TOKEN:
     raise ValueError("Hugging Face token not found. Please set the 'NEX_MODEL' environment variable.")
 
-login(token=HF_TOKEN)
+
+
 # ==============================
 # Helper: Human-readable bytes
 def sizeof_fmt(num, suffix="B"):
@@ -82,27 +81,37 @@ def calculate_cache_size(cache):
     return total_memory /(1024*1024)
     
 @st.cache_resource
-def load_model_and_tokenizer(doc_text_count):
-    model_name = "google/gemma-3-4b-it"  # Configure quantization for 4-bit loading
-  
-    # Load the pre-trained model with quantization
-    model = AutoModelForCausalLM.from_pretrained(
+def load_model_and_tokenizer():
+    model_name = "GeneZC/MiniChat-1.5-3B"
+        
+    
+    tokenizer = AutoTokenizer.from_pretrained(
         model_name,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto",
         trust_remote_code=True
         ,token=HF_TOKEN
     )
-
-    # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        trust_remote_code=True,
-        model_max_length=1.3*round(doc_text_count * 0.3 + 1)
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto",
+        trust_remote_code=True
         ,token=HF_TOKEN
     )
-    return tokenizer, model
-
+    return model, tokenizer
+def calculate_cache_size(cache):
+    """
+    Calculate the total memory used by the key-value cache (past_key_values) in megabytes.
+    Args:
+        cache: The past_key_values object (usually a tuple of (key, value) pairs per layer).
+    Returns:
+        Total memory in megabytes.
+    """
+    total_memory = 0
+    for layer_cache in cache:
+        key_tensor, value_tensor = layer_cache
+        total_memory += key_tensor.element_size() * key_tensor.nelement()
+        total_memory += value_tensor.element_size() * value_tensor.nelement()
+    return total_memory / (1024 * 1024)  # Convert to MB
 def clone_cache(cache):
     new_cache = DynamicCache()
     for key, value in zip(cache.key_cache, cache.value_cache):
@@ -117,7 +126,16 @@ def load_document_and_cache(file_path):
         with open(file_path, 'r') as file:
             doc_text = file.read()
         doc_text_count = len(doc_text)
-        model, tokenizer = load_model_and_tokenizer(doc_text_count)
+        max_length = int(1.3 * (doc_text_count * 0.3 + 1))
+    
+        # Cap the value at 16824
+        if max_length > 16824:
+            max_length = 16824
+        print(f" model_max_length set to: {max_length}")
+        
+        model, tokenizer = load_model_and_tokenizer()
+        tokenizer.model_max_length=max_length
+
         system_prompt = f"""
         <|system|>
         You are a helpful assistant. Provide concise, factual answers based only on the provided context.
@@ -194,11 +212,10 @@ if uploaded_file:
     print(f"👀 Document Preview Display Time: {t_end3 - t_start3:.2f} s")
     t_start4 = time()
     # PART 4: Show Basic Info
-    #doc_size_kb = os.path.getsize(temp_file_path) / 1024
-    #cache_size = os.path.getsize("temp_cache.pth") / 1024 if os.path.exists("temp_cache.pth") else "N/A"
+    s_cache=calculate_cache_size(cache)
     t_end4 = time()
     log.append(f"👀 doc_size_kb Preview Display Time: {t_end4 - t_start4:.2f} s")
-    print(f"👀 doc_size_kb Preview Display Time: {t_end4 - t_start4:.2f} s")
+    print(f"👀 doc_size_kb Preview Display Time: {t_end4 - t_start4:.2f} s||||||| size of the cache : {s_cache} MB")
     #st.info(
       #  f"Document Chars: {len(doc_text)} | Size: {doc_size_kb:.2f} KB | "
        # f"Cache Size: {cache_size if cache_size == 'N/A' else f'{cache_size:.2f} KB'}"
@@ -222,10 +239,10 @@ if uploaded_file:
 
             # PART 4.2: Tokenize Prompt
             t_start6 = time()
-            model, tokenizer = load_model_and_tokenizer(doc_text_count)
+
             full_prompt = f"""
             <|user|>
-            Question: {query}
+            Question: Please provide a clear and concise answer to the question .{query}
             <|assistant|>
             """.strip()
             input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids
@@ -249,7 +266,7 @@ if uploaded_file:
 
             st.success("Answer:")
             st.write(response)
-
+            print(f"***************************************************************************************")
             # Final Info Display
             st.info(
              #   f"Document Chars: {len(doc_text)} | Size: {doc_size_kb:.2f} KB | "