Update app.py

app.py

@@ -1,21 +1,23 @@
+import os
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
+from huggingface_hub import login
 
-# Load Hugging Face Token from environment variable
-hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+# Load Hugging Face Token from Secrets
+hf_token = os.getenv("HF_TOKEN")
 
-if hf_token is None:
-    st.error("❌ Hugging Face API Token is missing! Add it in Settings → Secrets.")
+if not hf_token:
+    st.error("Hugging Face token is missing! Please add it to Hugging Face Secrets.")
     st.stop()
 
-# Model name
-MODEL_NAME = "google/gemma-2b-it"
+# Authenticate
+login(token=hf_token)
 
-# Load tokenizer and model with authentication
-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)
+# Load the model and tokenizer with authentication
+MODEL_NAME = "google/gemma-2b-it"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=hf_token, torch_dtype=torch.float16, device_map="auto")
 
 # Streamlit UI
 st.title("Gemma-2B Code Assistant")