Update app.py

app.py

@@ -1,10 +1,16 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import login
 import torch
+import os
+
+# Use token from environment variable (set securely in Hugging Face Space)
+token = os.environ.get("HF_TOKEN")
+login(token)
 
 model_id = "bigcode/starcoder"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", use_auth_token=token)
 
 def generate_code(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -12,5 +18,4 @@ def generate_code(prompt):
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 iface = gr.Interface(fn=generate_code, inputs="text", outputs="text", title="StarCoder Code Generator")
-
 iface.launch(server_name="0.0.0.0", server_port=7860)