Update main.py

main.py

@@ -6,6 +6,14 @@ from pydantic import BaseModel
 from typing import Optional, Any
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer, GenerationConfig
+import os
+
+try:
+    # Attempt to fetch the token from the environment variable set in Spaces
+    HUGGINGFACE_TOKEN = os.environ['HF_ACCESS_TOKEN']
+except KeyError:
+    print('The environment variable "HF_ACCESS_TOKEN" is not found. Please configure it correctly in your Space.')
+    sys.exit(1)
 
 # Authentication
 from huggingface_hub import interpreter_login