Update app.py

app.py

@@ -1,14 +1,15 @@
 import gradio as gr
-import torch
+import torch, os
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import StoppingCriteria, TextIteratorStreamer
 from threading import Thread
 
 torch.set_num_threads(2)
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # Loading the tokenizer and model from Hugging Face's model hub.
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", use_auth_token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", use_auth_token=HF_TOKEN)
 
 def count_tokens(text):
     return len(tokenizer.tokenize(text))