Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,17 +1,22 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import streamlit as st
 import torch
-import os
-
 
 st.title("Tokenizer Test Space")
-
 model_id = "google/gemma-2b-it" # Test with the official model first
 # model_id = "Rahul-8799/project_manager_gemma3" # If the official model works, try yours
 
+# Define a writable directory for the cache. /tmp is usually writable in Spaces.
+cache_directory = "/tmp/hf_cache"
+
+# Ensure the cache directory exists (good practice, though hf_hub might handle it)
+os.makedirs(cache_directory, exist_ok=True)
+
 try:
     st.write(f"Attempting to load tokenizer for {model_id}...")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    # Explicitly pass the cache_dir
+    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_directory)
     st.success("Tokenizer loaded successfully!")
     st.write("Tokenizer details:", tokenizer)
 except Exception as e:
@@ -21,19 +26,20 @@ except Exception as e:
 try:
     st.write(f"Attempting to load model for {model_id}...")
     # Assuming you want 4-bit quantization for Gemma
-    from transformers import BitsAndBytesConfig
     quantization_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_quant_type="nf4",
         bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_use_double_quant=False,
     )
+    # Explicitly pass the cache_dir
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         quantization_config=quantization_config,
         low_cpu_mem_usage=True,
         torch_dtype=torch.bfloat16,
-        trust_remote_code=True
+        trust_remote_code=True,
+        cache_dir=cache_directory # Add this line
     )
     st.success("Model loaded successfully!")
     st.write("Model details:", model)