Update app.py

app.py

@@ -42,13 +42,21 @@ def get_labels(response_list):
 print(f"Starting to load the model to memory")
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
 m = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-2b-it", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True)
+    "google/gemma-2b-it", 
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, 
+    trust_remote_code=True,token=HF_TOKEN
+)
+
 embedding_func=m.get_input_embeddings()
 embedding_func.weight.requires_grad=False
 m = m.to(device)
 
-tok = AutoTokenizer.from_pretrained("google/gemma-2b-it", trust_remote_code=True)
+tok = AutoTokenizer.from_pretrained("google/gemma-2b-it", 
+                                    trust_remote_code=True,token=HF_TOKEN
+                                   )
 tok.padding_side = "left"
 tok.pad_token_id = tok.eos_token_id
 # using CUDA for an optimal experience