Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,26 +1,32 @@
 import json, re, ast, streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch 
+import os # Necessary to read the HF_TOKEN from environment variables
 
-# SWITCHED MODEL: From Mistral-7B to the much smaller Gemma-2B-Instruct
+# Model ID for the small, structured Gemma model
 model_id = "google/gemma-2b-it"
 
-tok = AutoTokenizer.from_pretrained(model_id)
+# Get the Hugging Face Token from the Space Secrets
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# Simplified Model Loading: Removed BitsAndBytesConfig
-# This smaller model might load cleanly without 4-bit quantization, resolving the dependency issues.
+tok = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN) 
+
+# Simplified Model Loading: No quantization needed due to smaller size
 try:
+    # Attempt to load using bfloat16 for efficiency
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
         torch_dtype=torch.bfloat16, 
-        device_map="auto"
+        device_map="auto",
+        token=HF_TOKEN 
     )
 except Exception:
-    # Fallback to float16 if bfloat16 causes issues
+    # Fallback to float16 if bfloat16 is not supported
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
         torch_dtype=torch.float16, 
-        device_map="auto"
+        device_map="auto",
+        token=HF_TOKEN
     )
 
 gen = pipeline("text-generation", model=model, tokenizer=tok,