Update app.py

app.py

@@ -1,22 +1,20 @@
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
 import os
 
-# Authenticate using the secret
-hf_token = os.getenv("HF_TOKEN")
-if hf_token:
-    login(token=hf_token)
-else:
-    print("Hugging Face token not found! Please add it to Secrets.")
-
-
+# Load Hugging Face token from environment variable (Set this in Hugging Face Secrets)
+hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+if not hf_token:
+    st.error("Hugging Face token not found! Please add it to Secrets in Hugging Face Spaces.")
+    st.stop()
 
 # Load the model and tokenizer
 MODEL_NAME = "google/gemma-2b-it"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, torch_dtype=torch.float16, device_map="auto", use_auth_token=hf_token
+)
 
 # Streamlit UI
 st.title("Gemma-2B Code Assistant")